commit f1dd9344fe4c5ebbcd5a0901ba206307edd490d5
Author: wehub-resource-sync <resource-sync@wehub.local>
Date:   Mon Jul 13 13:16:24 2026 +0800

    chore: import upstream snapshot with attribution

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..000fa71
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,26 @@
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+
+# Output files
+output.glb
+autotune_cache.json
+
+# Virtual environment
+venv/
+.venv/
+env/
+
+# IDE
+.idea/
+.vscode/
+*.swp
+*.swo
+
+# OS
+.DS_Store
+Thumbs.db
+
+tmp
diff --git a/LICENSE b/LICENSE
new file mode 100644
index 0000000..9e907dc
--- /dev/null
+++ b/LICENSE
@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Tencent.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
diff --git a/NOTICE b/NOTICE
new file mode 100644
index 0000000..cab8abe
--- /dev/null
+++ b/NOTICE
@@ -0,0 +1,132 @@
+Pixal3D
+Copyright (c) 2026 Tencent.
+
+Tencent is pleased to support the community by making Pixal3D available.
+
+The open-source software and/or models included in this distribution may have
+been modified by Tencent (“Tencent Modifications”). All Tencent Modifications
+are Copyright (c) Tencent.
+
+For the avoidance of doubt, Pixal3D refers solely to code, parameters, weights,
+and associated documentation made publicly available by Tencent in accordance
+with the MIT License (see the LICENSE file).
+
+Responsible Use
+This model is released to support research, creative production, and practical 3D generation. Users are responsible for ensuring that their use complies with applicable laws, rights, consent requirements, and platform policies.
+
+We do not support use of this model to create non-consensual sexual content, sexualized depictions of real people, content involving minors, harassment, impersonation, or privacy-violating applications.
+
+The model is provided “as is” under the MIT License. This guidance does not modify the license terms but describes the intended and responsible use of the project.
+
+================================================================================
+Third-Party Components
+================================================================================
+
+This product includes software developed by third parties. The third-party
+components listed below remain licensed under their respective original terms.
+Pixal3D does not impose any additional restrictions beyond those specified in
+the original licenses of these third-party components. Users are required to
+comply with all applicable terms and conditions of the original licenses and
+to ensure that the use of these third-party components conforms to all
+relevant laws and regulations.
+
+In case you believe there have been errors in the attribution below, you may
+submit the concerns to us for review and correction.
+
+
+================================================================================
+Open Model Licensed under Apache-2.0
+================================================================================
+
+1. dinov2
+   Copyright (c) 2025 dinov2 original author and authors
+
+
+================================================================================
+Open Source Software Licensed under the MIT License
+================================================================================
+
+2. TRELLIS.2
+   Copyright (c) Microsoft Corporation.
+
+3. Direct3D-S2
+   Copyright (c) DreamTech.
+
+4. MoGe
+   Copyright (c) Microsoft Corporation.
+
+
+================================================================================
+Terms of the MIT License
+================================================================================
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+
+================================================================================
+Terms of the Apache License, Version 2.0
+================================================================================
+
+Apache License
+                       Version 2.0, January 2004
+                    http://www.apache.org/licenses/
+
+TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+Definitions.
+
+"License" shall mean the terms and conditions for use, reproduction, and distribution as defined by Sections 1 through 9 of this document.
+
+"Licensor" shall mean the copyright owner or entity authorized by the copyright owner that is granting the License.
+
+"Legal Entity" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity.
+
+"You" (or "Your") shall mean an individual or Legal Entity exercising permissions granted by this License.
+
+"Source" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files.
+
+"Object" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types.
+
+"Work" shall mean the work of authorship, whether in Source or Object form, made available under the License, as indicated by a copyright notice that is included in or attached to the work (an example is provided in the Appendix below).
+
+"Derivative Works" shall mean any work, whether in Source or Object form, that is based on (or derived from) the Work and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this License, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work and Derivative Works thereof.
+
+"Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally submitted to Licensor for inclusion in the Work by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright 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 and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution."
+
+"Contributor" shall mean Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Work.
+
+Grant of Copyright License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare Derivative Works of, publicly display, publicly perform, sublicense, and distribute the Work and such Derivative Works in Source or Object form.
+
+Grant of Patent License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Work to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or a Contribution incorporated within the Work constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed.
+
+Redistribution. You may reproduce and distribute copies of the Work or Derivative Works thereof in any medium, with or without modifications, and in Source or Object form, provided that You meet the following conditions:
+
+(a) You must give any other recipients of the Work or Derivative Works a copy of this License; and (b) You must cause any modified files to carry prominent notices stating that You changed the files; and (c) You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works; and (d) If the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation, if provided along with the Derivative Works; or, within a display generated by the Derivative Works, if and wherever such third-party notices normally appear. The contents of the NOTICE file are for informational purposes only and do not modify the License. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from the Work, provided that such additional attribution notices cannot be construed as modifying the License. You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Works as a whole, provided Your use, reproduction, and distribution of the Work otherwise complies with the conditions stated in this License.
+
+Submission of Contributions. Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions.
+
+Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor, except as required for reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file.
+
+Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Work (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 or redistributing the Work and assume any risks associated with Your exercise of permissions under this License.
+
+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 Work (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.
+
+Accepting Warranty or Additional Liability. While redistributing the Work or Derivative Works thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability.
+
+END OF TERMS AND CONDITIONS
diff --git a/README.md b/README.md
new file mode 100644
index 0000000..09d7884
--- /dev/null
+++ b/README.md
@@ -0,0 +1,298 @@
+
+<div align="center">
+
+# Pixal3D: Pixel-Aligned 3D Generation from Images
+
+<h3>SIGGRAPH 2026</h3>
+
+<small>[Dong-Yang Li](https://ldyang694.github.io/)¹ · [Wang Zhao](https://thuzhaowang.github.io/)²* · [Yuxin Chen](https://orcid.org/0000-0002-7854-1072)² · [Wenbo Hu](https://wbhu.github.io/)² · [Meng-Hao Guo](https://menghaoguo.github.io/)¹ · [Fang-Lue Zhang](https://fanglue.github.io/)³ · [Ying Shan](https://www.linkedin.com/in/YingShanProfile)² · [Shi-Min Hu](https://cg.cs.tsinghua.edu.cn/shimin.htm)¹✉</small>
+
+¹Tsinghua University (BNRist) &nbsp;&nbsp; ²Tencent ARC Lab &nbsp;&nbsp; ³Victoria University of Wellington
+
+*Project lead &nbsp;&nbsp; ✉Corresponding author
+
+</div>
+
+<div align="center">
+  <a href="https://ldyang694.github.io/projects/pixal3d/"><img src=https://img.shields.io/badge/Project%20Page-333399.svg?logo=googlehome height=22px></a>
+  <a href="https://huggingface.co/spaces/TencentARC/Pixal3D"><img src=https://img.shields.io/badge/%F0%9F%A4%97%20Demo-276cb4.svg height=22px></a>
+  <a href="https://huggingface.co/TencentARC/Pixal3D"><img src=https://img.shields.io/badge/%F0%9F%A4%97%20Models-d96902.svg height=22px></a>
+  <a href="https://arxiv.org/abs/2605.10922"><img src=https://img.shields.io/badge/Arxiv-b5212f.svg?logo=arxiv height=22px></a>
+  <a href="LICENSE"><img src=https://img.shields.io/badge/License-MIT-yellow.svg height=22px></a>
+</div>
+
+<div align="center">
+    <img src="assets/teaser.png" alt="Teaser image of Pixal3D"/>
+</div>
+
+**Pixal3D** generates high-fidelity 3D assets from a single image. Unlike previous methods that loosely inject image features via attention, Pixal3D explicitly lifts pixel features into 3D through back-projection, establishing direct pixel-to-3D correspondences. This enables near-reconstruction-level fidelity with detailed geometry and PBR textures.
+
+---
+
+## ✨ News
+
+- **May 2026**: Release training code and data preparation toolkit. 🔧
+- **May 2026**: Release the improved version based on [Trellis.2](https://github.com/microsoft/TRELLIS.2) backbone. 💪
+- **May 2026**: Release inference code and online demo. 🤗
+- **Apr 2026**: Our paper is accepted to SIGGRAPH 2026! 🎉
+
+## 📌 Branches
+
+| Branch | Description |
+|--------|-------------|
+| `main` | **Latest version** — improved implementation based on [Trellis.2](https://github.com/microsoft/TRELLIS.2) backbone with better performance. |
+| `paper` | **Paper version** — original implementation based on [Direct3D-S2](https://github.com/DreamTechAI/Direct3D-S2), corresponding to results reported in our SIGGRAPH 2026 paper. |
+
+> If you want to reproduce the results in our paper, please switch to the `paper` branch.
+
+## 🎮 Try It Online
+
+You can try Pixal3D directly in your browser without any installation via our Hugging Face Gradio demo:
+
+👉 [**Launch Demo**](https://huggingface.co/spaces/TencentARC/Pixal3D)
+
+## 🚀 Getting Started
+
+### Installation
+
+#### Step 1: Follow TRELLIS.2 Installation
+
+Please first follow the installation guide of [TRELLIS.2](https://github.com/microsoft/TRELLIS.2) to set up the base environment.
+
+#### Step 2: Install Additional Dependencies
+
+```bash
+pip install -r requirements.txt
+```
+
+#### Step 3: Install natten
+
+```bash
+NATTEN_CUDA_ARCH="xx" NATTEN_N_WORKERS=xx pip install natten==0.21.0 --no-build-isolation
+```
+
+Please replace `xx` with the CUDA architecture and the number of build workers suitable for your machine.
+
+#### Step 4: Install utils3d
+
+```bash
+pip install https://github.com/LDYang694/Storages/releases/download/20260430/utils3d-0.0.2-py3-none-any.whl
+```
+
+> **Note**: `requirements-hfdemo.txt` is for the Hugging Face Spaces demo (H-series GPU architecture) and may not be compatible with other architectures.
+
+### Usage
+
+#### Inference
+
+Generate a GLB mesh from a single image:
+
+```bash
+python inference.py --image assets/images/0_img.png --output ./output.glb
+```
+
+**Low-VRAM mode** (reduces peak VRAM by loading models on-demand):
+
+```bash
+python inference.py --image assets/images/0_img.png --output ./output.glb --low_vram
+```
+
+By default, the pipeline resolution is **1536** (standard mode) or **1024** (low-VRAM mode). You can override this with `--resolution`:
+
+```bash
+# Force 1536 even in low-VRAM mode
+python inference.py --image assets/images/0_img.png --output ./output.glb --low_vram --resolution 1536
+
+# Force 1024 in standard mode
+python inference.py --image assets/images/0_img.png --output ./output.glb --resolution 1024
+```
+
+**Tip**: If you don't have `flash_attn` installed, you can use PyTorch's built-in SDPA backend instead:
+> ```bash
+> ATTN_BACKEND=sdpa python inference.py --image assets/images/0_img.png --output ./output.glb --low_vram
+> ```
+
+### Web Demo
+
+We provide a Gradio web demo for Pixal3D, which allows you to generate 3D meshes from images interactively.
+
+```bash
+python app.py 
+```
+
+Low-VRAM mode is also available for the web demo. The frontend default resolution will automatically switch to 1024 in low-VRAM mode (1536 otherwise), but can be changed manually in the UI.
+
+```bash
+python app.py --low_vram
+# or via environment variable:
+LOW_VRAM=1 python app.py
+```
+## 🔧 Training
+
+We provide the full training codebase for reproducing Pixal3D from scratch.
+
+### Data Preparation
+
+Prepare view-aligned O-Voxel data and rendered condition images by following the data toolkit instructions:
+
+> 📂 **[data_toolkit/README.md](data_toolkit/README.md)**
+
+### Overview
+
+Pixal3D is trained as a three-stage cascade, each progressively increasing resolution:
+
+| Stage | Model | Resolutions | Config Prefix |
+|-------|-------|-------------|---------------|
+| 1 | Sparse Structure | 32 → 64 | `ss_flow_img_dit_*_proj_finetune` |
+| 2 | Shape | 256 → 512 → 1024 | `slat_flow_img2shape_*_proj_finetune` |
+| 3 | Texture | 256 → 512 → 1024 | `slat_flow_imgshape2tex_*_proj_finetune` |
+
+All stages use **pixel-aligned projection conditioning** and **view-aligned latents** (2 views by default). Within each stage, start from the lowest resolution and progressively fine-tune to higher resolutions by setting `finetune_ckpt` in the config.
+
+### Quick Start
+
+```sh
+python train.py \
+  --config <CONFIG_JSON> \
+  --output_dir <OUTPUT_DIR> \
+  --data_dir '<DATA_DIR_JSON>'
+```
+
+`--data_dir` is a JSON string describing the dataset layout. Different stages require different keys:
+
+| Stage | Required keys |
+|-------|---------------|
+| Sparse Structure | `base`, `ss_latent`, `render_cond` |
+| Shape | `base`, `shape_latent`, `render_cond` |
+| Texture | `base`, `shape_latent`, `pbr_latent`, `render_cond` |
+
+### Example: Training All Three Stages
+
+Below we show the full training sequence using ObjaverseXL as an example. Each higher-resolution step requires updating `finetune_ckpt` in its config JSON to point to the previous checkpoint.
+
+<details>
+<summary><b>Stage 1: Sparse Structure (32 → 64)</b></summary>
+
+```sh
+# Resolution 32
+python train.py \
+  --config configs/gen/ss_flow_img_dit_1_3B_32_bf16_proj_finetune.json \
+  --output_dir results/ss_32 \
+  --data_dir '{"ObjaverseXL_sketchfab": {"base": "datasets/ObjaverseXL_sketchfab", "ss_latent": "datasets/ObjaverseXL_sketchfab/ss_latents/ss_enc_conv3d_16l8_fp16_64_view", "render_cond": "datasets/ObjaverseXL_sketchfab/renders_cond"}}'
+
+# Resolution 64 (set finetune_ckpt → results/ss_32 checkpoint)
+python train.py \
+  --config configs/gen/ss_flow_img_dit_1_3B_32_bf16_proj_finetune_ft64.json \
+  --output_dir results/ss_ft64 \
+  --data_dir '{"ObjaverseXL_sketchfab": {"base": "datasets/ObjaverseXL_sketchfab", "ss_latent": "datasets/ObjaverseXL_sketchfab/ss_latents/ss_enc_conv3d_16l8_fp16_64_view", "render_cond": "datasets/ObjaverseXL_sketchfab/renders_cond"}}'
+```
+</details>
+
+<details>
+<summary><b>Stage 2: Shape (256 → 512 → 1024)</b></summary>
+
+```sh
+# Resolution 256
+python train.py \
+  --config configs/gen/slat_flow_img2shape_dit_1_3B_256_bf16_proj_finetune.json \
+  --output_dir results/shape_256 \
+  --data_dir '{"ObjaverseXL_sketchfab": {"base": "datasets/ObjaverseXL_sketchfab", "shape_latent": "datasets/ObjaverseXL_sketchfab/shape_latents/shape_enc_next_dc_f16c32_fp16_256_view", "render_cond": "datasets/ObjaverseXL_sketchfab/renders_cond"}}'
+
+# Resolution 512
+python train.py \
+  --config configs/gen/slat_flow_img2shape_dit_1_3B_256_bf16_proj_finetune_ft512.json \
+  --output_dir results/shape_ft512 \
+  --data_dir '{"ObjaverseXL_sketchfab": {"base": "datasets/ObjaverseXL_sketchfab", "shape_latent": "datasets/ObjaverseXL_sketchfab/shape_latents/shape_enc_next_dc_f16c32_fp16_512_view", "render_cond": "datasets/ObjaverseXL_sketchfab/renders_cond"}}'
+
+# Resolution 1024
+python train.py \
+  --config configs/gen/slat_flow_img2shape_dit_1_3B_512_bf16_proj_finetune_ft1024.json \
+  --output_dir results/shape_ft1024 \
+  --data_dir '{"ObjaverseXL_sketchfab": {"base": "datasets/ObjaverseXL_sketchfab", "shape_latent": "datasets/ObjaverseXL_sketchfab/shape_latents/shape_enc_next_dc_f16c32_fp16_1024_view", "render_cond": "datasets/ObjaverseXL_sketchfab/renders_cond"}}'
+```
+</details>
+
+<details>
+<summary><b>Stage 3: Texture (256 → 512 → 1024)</b></summary>
+
+```sh
+# Resolution 256
+python train.py \
+  --config configs/gen/slat_flow_imgshape2tex_dit_1_3B_256_bf16_proj_finetune.json \
+  --output_dir results/tex_256 \
+  --data_dir '{"ObjaverseXL_sketchfab": {"base": "datasets/ObjaverseXL_sketchfab", "shape_latent": "datasets/ObjaverseXL_sketchfab/shape_latents/shape_enc_next_dc_f16c32_fp16_256_view", "pbr_latent": "datasets/ObjaverseXL_sketchfab/pbr_latents/tex_enc_next_dc_f16c32_fp16_256_view", "render_cond": "datasets/ObjaverseXL_sketchfab/renders_cond"}}'
+
+# Resolution 512
+python train.py \
+  --config configs/gen/slat_flow_imgshape2tex_dit_1_3B_512_bf16_proj_finetune.json \
+  --output_dir results/tex_512 \
+  --data_dir '{"ObjaverseXL_sketchfab": {"base": "datasets/ObjaverseXL_sketchfab", "shape_latent": "datasets/ObjaverseXL_sketchfab/shape_latents/shape_enc_next_dc_f16c32_fp16_512_view", "pbr_latent": "datasets/ObjaverseXL_sketchfab/pbr_latents/tex_enc_next_dc_f16c32_fp16_512_view", "render_cond": "datasets/ObjaverseXL_sketchfab/renders_cond"}}'
+
+# Resolution 1024
+python train.py \
+  --config configs/gen/slat_flow_imgshape2tex_dit_1_3B_512_bf16_proj_finetune_ft1024.json \
+  --output_dir results/tex_ft1024 \
+  --data_dir '{"ObjaverseXL_sketchfab": {"base": "datasets/ObjaverseXL_sketchfab", "shape_latent": "datasets/ObjaverseXL_sketchfab/shape_latents/shape_enc_next_dc_f16c32_fp16_1024_view", "pbr_latent": "datasets/ObjaverseXL_sketchfab/pbr_latents/tex_enc_next_dc_f16c32_fp16_1024_view", "render_cond": "datasets/ObjaverseXL_sketchfab/renders_cond"}}'
+```
+</details>
+
+### Additional Options
+
+<details>
+<summary><b>All command-line arguments</b></summary>
+
+| Argument | Description | Default |
+|----------|-------------|---------|
+| `--config` | Config JSON path | *required* |
+| `--output_dir` | Output directory | *required* |
+| `--data_dir` | Dataset JSON string | `./data/` |
+| `--load_dir` | Checkpoint load directory | `output_dir` |
+| `--ckpt` | Resume from step | `latest` |
+| `--auto_retry` | Retries on failure | `3` |
+| `--tryrun` | Dry run | `false` |
+| `--profile` | Profiling | `false` |
+| `--num_nodes` | Number of nodes | `1` |
+| `--node_rank` | Current node rank | `0` |
+| `--num_gpus` | GPUs per node | all |
+| `--master_addr` | Master address | `localhost` |
+| `--master_port` | Master port | `12666` |
+| `--use_wandb` | Enable W&B logging | `false` |
+| `--wandb_project` | W&B project | `trellis2-training` |
+| `--wandb_name` | W&B run name | basename of `output_dir` |
+| `--wandb_id` | W&B run ID (resume) | — |
+
+</details>
+
+## 🌐 Community Projects
+
+We thank the community for building extensions and deployment guides for Pixal3D!
+
+- [Pixal3D-ComfyUI](https://github.com/Saganaki22/Pixal3D-ComfyUI) — ComfyUI integration with deployment guides for Windows, WSL, and more.
+
+## 🤗 Acknowledgements
+
+This project is heavily built upon [Trellis.2](https://github.com/microsoft/TRELLIS.2) and [Direct3D-S2](https://github.com/DreamTechAI/Direct3D-S2). We sincerely thank the authors for their outstanding work on scalable 3D generation , which serves as the foundation of our codebase and model architecture.
+
+We also thank the following repos for their great contributions:
+
+- [Direct3D-S2](https://github.com/DreamTechAI/Direct3D-S2)
+- [Trellis](https://github.com/microsoft/TRELLIS)
+- [Trellis.2](https://github.com/microsoft/TRELLIS.2)
+
+## 📄 Citation
+
+If you find this work useful, please consider citing:
+
+```bibtex
+@article{li2026pixal3d,
+    title={Pixal3D: Pixel-Aligned 3D Generation from Images},
+    author={Li, Dong-Yang and Zhao, Wang and Chen, Yuxin and Hu, Wenbo and Guo, Meng-Hao and Zhang, Fang-Lue and Shan, Ying and Hu, Shi-Min},
+    journal={arXiv preprint arXiv:2605.10922},
+    year={2026}
+}
+```
+
+## 📜 License
+
+This project is released under the [MIT License](LICENSE). The third-party components included in this project remain licensed under their respective original terms; see [NOTICE](NOTICE) for the full list of dependencies and their licenses.
+
diff --git a/README.wehub.md b/README.wehub.md
new file mode 100644
index 0000000..6e5113b
--- /dev/null
+++ b/README.wehub.md
@@ -0,0 +1,7 @@
+# WeHub 来源说明
+
+- 原始项目：`TencentARC/Pixal3D`
+- 原始仓库：https://github.com/TencentARC/Pixal3D
+- 导入方式：上游默认分支的最新快照
+- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准
+- 本文件仅用于记录来源，不代表 WeHub 是原项目作者
diff --git a/app.py b/app.py
new file mode 100644
index 0000000..48479df
--- /dev/null
+++ b/app.py
@@ -0,0 +1,559 @@
+import os
+import subprocess
+import argparse
+import math
+import time
+import shutil
+import cv2
+import torch
+import numpy as np
+import base64
+import io
+import json
+from datetime import datetime
+from typing import *
+from PIL import Image
+
+import threading
+try:
+    import nest_asyncio
+    nest_asyncio.apply()
+except ImportError:
+    pass
+
+# Lock for model initialization
+init_lock = threading.Lock()
+
+os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1'
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+os.environ.setdefault("ATTN_BACKEND", "flash_attn")
+os.environ["FLEX_GEMM_AUTOTUNE_CACHE_PATH"] = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'autotune_cache.json')
+os.environ["FLEX_GEMM_AUTOTUNER_VERBOSE"] = '1'
+
+import spaces
+from gradio import Server
+from gradio.data_classes import FileData
+from fastapi.responses import HTMLResponse
+from fastapi.staticfiles import StaticFiles
+
+from pixal3d.modules.sparse import SparseTensor
+from pixal3d.pipelines import Pixal3DImageTo3DPipeline
+from pixal3d.renderers import EnvMap
+from pixal3d.utils import render_utils
+import o_voxel
+
+# ============================================================================
+# Constants & Defaults
+# ============================================================================
+
+MAX_SEED = np.iinfo(np.int32).max
+TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
+os.makedirs(TMP_DIR, exist_ok=True)
+
+MODES = [
+    {"name": "Normal", "icon": "assets/app/normal.png", "render_key": "normal"},
+    {"name": "Clay render", "icon": "assets/app/clay.png", "render_key": "clay"},
+    {"name": "Base color", "icon": "assets/app/basecolor.png", "render_key": "base_color"},
+    {"name": "HDRI forest", "icon": "assets/app/hdri_forest.png", "render_key": "shaded_forest"},
+    {"name": "HDRI sunset", "icon": "assets/app/hdri_sunset.png", "render_key": "shaded_sunset"},
+    {"name": "HDRI courtyard", "icon": "assets/app/hdri_courtyard.png", "render_key": "shaded_courtyard"},
+]
+STEPS = 8
+
+# Cascade parameters
+CASCADE_LR_RESOLUTION = 512
+CASCADE_MAX_NUM_TOKENS = 49152
+
+# MoGe defaults
+MOGE_MODEL_NAME = "Ruicheng/moge-2-vitl"
+WILD_MESH_SCALE = 1.0
+WILD_EXTEND_PIXEL = 0
+WILD_IMAGE_RESOLUTION = 512
+
+# Image Cond Model configs
+IMAGE_COND_CONFIGS = {
+    "ss": {
+        "model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
+        "image_size": 512,
+        "grid_resolution": 16,
+    },
+    "shape_512": {
+        "model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
+        "image_size": 512,
+        "grid_resolution": 32,
+        "use_naf_upsample": True,
+        "naf_target_size": 512,
+    },
+    "shape_1024": {
+        "model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
+        "image_size": 1024,
+        "grid_resolution": 64,
+        "use_naf_upsample": True,
+        "naf_target_size": 512,
+    },
+    "tex_1024": {
+        "model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
+        "image_size": 1024,
+        "grid_resolution": 64,
+        "use_naf_upsample": True,
+        "naf_target_size": 1024,
+    },
+}
+
+# ============================================================================
+# Model Loading
+# ============================================================================
+
+def build_image_cond_model(config: dict):
+    from pixal3d.trainers.flow_matching.mixins.image_conditioned_proj import DinoV3ProjFeatureExtractor
+    model = DinoV3ProjFeatureExtractor(**config)
+    model.eval()
+    return model
+
+def load_moge_model(device="cuda", model_name=MOGE_MODEL_NAME):
+    from moge.model.v2 import MoGeModel
+    moge_model = MoGeModel.from_pretrained(model_name).to(device)
+    moge_model.eval()
+    return moge_model
+
+# Global instances (lazy loaded or loaded at start)
+pipeline = None
+moge_model = None
+envmap = None
+LOW_VRAM = os.environ.get("LOW_VRAM", "0") == "1"
+
+def init_models():
+    global pipeline, moge_model, envmap
+    with init_lock:
+        if pipeline is not None:
+            return
+
+        # GPU / CUDA Diagnostics (runs when GPU is allocated)
+        import subprocess as _sp
+        print("=" * 60)
+        print("[Diagnostics] PyTorch version:", torch.__version__)
+        print("[Diagnostics] CUDA available:", torch.cuda.is_available())
+        if torch.cuda.is_available():
+            print("[Diagnostics] CUDA version:", torch.version.cuda)
+            print("[Diagnostics] cuDNN version:", torch.backends.cudnn.version())
+            for i in range(torch.cuda.device_count()):
+                name = torch.cuda.get_device_name(i)
+                cap = torch.cuda.get_device_capability(i)
+                mem = torch.cuda.get_device_properties(i).total_memory / 1024**3
+                print(f"[Diagnostics] GPU {i}: {name}, sm_{cap[0]}{cap[1]}, {mem:.1f} GB")
+        try:
+            res = _sp.run(["nvidia-smi", "--query-gpu=name,compute_cap,memory.total", "--format=csv,noheader"], capture_output=True, text=True, timeout=10)
+            print("[Diagnostics] nvidia-smi:", res.stdout.strip())
+        except Exception as e:
+            print(f"[Diagnostics] nvidia-smi failed: {e}")
+        print("=" * 60)
+
+        model_path = "TencentARC/Pixal3D"
+        print(f"[Pipeline] Loading from {model_path}...")
+        pipeline = Pixal3DImageTo3DPipeline.from_pretrained(model_path)
+        
+        print("[ImageCond] Building DinoV3ProjFeatureExtractor models...")
+        pipeline.image_cond_model_ss = build_image_cond_model(IMAGE_COND_CONFIGS["ss"])
+        pipeline.image_cond_model_shape_512 = build_image_cond_model(IMAGE_COND_CONFIGS["shape_512"])
+        pipeline.image_cond_model_shape_1024 = build_image_cond_model(IMAGE_COND_CONFIGS["shape_1024"])
+        pipeline.image_cond_model_tex_1024 = build_image_cond_model(IMAGE_COND_CONFIGS["tex_1024"])
+        
+        if LOW_VRAM:
+            # Low-VRAM mode: models stay on CPU, loaded to GPU on-demand per stage.
+            print("[NAF] Pre-downloading NAF upsampler weights (CPU only)...")
+            for attr in ['image_cond_model_ss', 'image_cond_model_shape_512',
+                         'image_cond_model_shape_1024', 'image_cond_model_tex_1024']:
+                m = getattr(pipeline, attr, None)
+                if m is not None and getattr(m, 'use_naf_upsample', False):
+                    m._load_naf()
+            pipeline._device = torch.device("cuda")
+            pipeline.low_vram = True
+            print("[Pipeline] Low-VRAM mode enabled.")
+        else:
+            # Standard mode: all models loaded to GPU at once.
+            pipeline.low_vram = False
+            pipeline.cuda()
+            pipeline.image_cond_model_ss.cuda()
+            pipeline.image_cond_model_shape_512.cuda()
+            pipeline.image_cond_model_shape_1024.cuda()
+            pipeline.image_cond_model_tex_1024.cuda()
+            print("[NAF] Pre-loading NAF upsampler model...")
+            for attr in ['image_cond_model_ss', 'image_cond_model_shape_512',
+                         'image_cond_model_shape_1024', 'image_cond_model_tex_1024']:
+                m = getattr(pipeline, attr, None)
+                if m is not None and getattr(m, 'use_naf_upsample', False):
+                    m._load_naf()
+                
+        print("[MoGe-2] Loading model for camera estimation...")
+        if LOW_VRAM:
+            # Low-VRAM: load MoGe to CPU, move to GPU on-demand per request.
+            moge_model = load_moge_model(device="cpu")
+            print("[MoGe-2] Low-VRAM mode: MoGe stays on CPU, loaded to GPU on-demand.")
+        else:
+            moge_model = load_moge_model(device="cuda")
+        
+        print("[EnvMap] Loading environment maps...")
+        _base = os.path.dirname(os.path.abspath(__file__))
+        _envmap_device = 'cpu' if LOW_VRAM else 'cuda'
+        envmap = {
+            'forest': EnvMap(torch.tensor(cv2.cvtColor(cv2.imread(os.path.join(_base, 'assets/hdri/forest.exr'), cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB), dtype=torch.float32, device=_envmap_device)),
+            'sunset': EnvMap(torch.tensor(cv2.cvtColor(cv2.imread(os.path.join(_base, 'assets/hdri/sunset.exr'), cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB), dtype=torch.float32, device=_envmap_device)),
+            'courtyard': EnvMap(torch.tensor(cv2.cvtColor(cv2.imread(os.path.join(_base, 'assets/hdri/courtyard.exr'), cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB), dtype=torch.float32, device=_envmap_device)),
+        }
+
+# ============================================================================
+# Utilities
+# ============================================================================
+
+def compute_f_pixels(camera_angle_x: float, resolution: int) -> float:
+    focal_length = 16.0 / torch.tan(torch.tensor(camera_angle_x / 2.0))
+    f_pixels = focal_length * resolution / 32.0
+    return float(f_pixels.item())
+
+def distance_from_fov(camera_angle_x, grid_point, target_point, mesh_scale, image_resolution):
+    rotation_matrix = torch.tensor([[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]])
+    gp = grid_point.to(torch.float32) @ rotation_matrix.T
+    gp = gp / mesh_scale / 2
+    xw, yw, zw = gp[0].item(), gp[1].item(), gp[2].item()
+    xt, yt = float(target_point[0].item()), float(target_point[1].item())
+    f_pixels = compute_f_pixels(camera_angle_x, image_resolution)
+    x_ndc = xt - image_resolution / 2.0
+    y_ndc = -(yt - image_resolution / 2.0)
+    distance_x = f_pixels * xw / x_ndc - yw
+    return {"distance_from_x": float(distance_x), "f_pixels": float(f_pixels)}
+
+def get_camera_params_wild_moge(image_path, device="cuda", mesh_scale=1.0, extend_pixel=0, image_resolution=512):
+    pil_image = Image.open(image_path).convert("RGB")
+    width, height = pil_image.size
+    image_np = np.array(pil_image).astype(np.float32) / 255.0
+    image_tensor = torch.from_numpy(image_np).permute(2, 0, 1).to(device)
+    if LOW_VRAM:
+        moge_model.to(device)
+    with torch.no_grad():
+        output = moge_model.infer(image_tensor)
+    if LOW_VRAM:
+        moge_model.cpu()
+        torch.cuda.empty_cache()
+    intrinsics = output["intrinsics"].squeeze().cpu().numpy()
+    fx_normalized = intrinsics[0, 0]
+    fx = fx_normalized * width
+    camera_angle_x = 2 * math.atan(width / (2 * fx))
+
+    grid_point = torch.tensor([-1.0, 0.0, 0.0])
+    distance = distance_from_fov(
+        camera_angle_x, grid_point,
+        torch.tensor([0 - extend_pixel, image_resolution - 1 + extend_pixel]),
+        mesh_scale, image_resolution
+    )["distance_from_x"]
+    return {'camera_angle_x': camera_angle_x, 'distance': distance, 'mesh_scale': mesh_scale}
+
+def pack_state(shape_slat, tex_slat, res):
+    state_data = {
+        'shape_slat_feats': shape_slat.feats.cpu().numpy(),
+        'tex_slat_feats': tex_slat.feats.cpu().numpy(),
+        'coords': shape_slat.coords.cpu().numpy(),
+        'res': res,
+    }
+    import random
+    state_path = os.path.join(TMP_DIR, f"state_{int(time.time()*1000)}_{random.randint(0,9999):04d}.npz")
+    np.savez_compressed(state_path, **state_data)
+    return state_path
+
+def unpack_state(state_path):
+    data = np.load(state_path)
+    shape_slat = SparseTensor(
+        feats=torch.from_numpy(data['shape_slat_feats']).cuda(),
+        coords=torch.from_numpy(data['coords']).cuda(),
+    )
+    tex_slat = shape_slat.replace(torch.from_numpy(data['tex_slat_feats']).cuda())
+    return shape_slat, tex_slat, int(data['res'])
+
+# ============================================================================
+# Progress Tracking (file-based, cross-process safe for @spaces.GPU)
+# ============================================================================
+
+import asyncio
+from fastapi.responses import JSONResponse
+from fastapi import Request
+
+PROGRESS_DIR = os.path.join(TMP_DIR, '_progress')
+os.makedirs(PROGRESS_DIR, exist_ok=True)
+
+_thread_local = threading.local()
+
+def _progress_file(session_id: str) -> str:
+    """Return path to a session's progress JSON file."""
+    return os.path.join(PROGRESS_DIR, f"{session_id}.json")
+
+def _reset_progress(session_id: str):
+    _thread_local.active_session = session_id
+    _write_progress_file(session_id, {"stage": "Initializing...", "step": 0, "total": 0, "done": False})
+
+def _update_progress(stage: str, step: int, total: int):
+    session_id = getattr(_thread_local, 'active_session', '')
+    if session_id:
+        _write_progress_file(session_id, {"stage": stage, "step": step, "total": total, "done": False})
+
+def _finish_progress():
+    session_id = getattr(_thread_local, 'active_session', '')
+    if session_id:
+        _write_progress_file(session_id, {"done": True})
+
+def _write_progress_file(session_id: str, data: dict):
+    """Atomically write progress JSON to a file (cross-process safe)."""
+    path = _progress_file(session_id)
+    tmp_path = path + ".tmp"
+    try:
+        with open(tmp_path, 'w') as f:
+            json.dump(data, f)
+        os.replace(tmp_path, path)  # atomic on POSIX
+    except Exception:
+        pass
+
+# Monkey-patch tqdm to intercept progress
+import tqdm as _tqdm_module
+
+_original_tqdm = _tqdm_module.tqdm
+
+class _TqdmProgressInterceptor(_original_tqdm):
+    """Wraps tqdm to push progress updates to SSE."""
+    def __init__(self, *args, **kwargs):
+        self._stage_desc = kwargs.get('desc', 'Processing')
+        super().__init__(*args, **kwargs)
+    
+    def set_description(self, desc=None, refresh=True):
+        self._stage_desc = desc or 'Processing'
+        super().set_description(desc, refresh)
+    
+    def update(self, n=1):
+        super().update(n)
+        _update_progress(self._stage_desc, self.n, self.total or 0)
+
+# Patch tqdm globally
+_tqdm_module.tqdm = _TqdmProgressInterceptor
+# Also patch the direct import in the sampler module and render_utils
+import pixal3d.pipelines.samplers.flow_euler as _fe_module
+_fe_module.tqdm = _TqdmProgressInterceptor
+import pixal3d.utils.render_utils as _ru_module
+_ru_module.tqdm = _TqdmProgressInterceptor
+import o_voxel.postprocess as _ovp_module
+_ovp_module.tqdm = _TqdmProgressInterceptor
+
+# ============================================================================
+# API Implementation
+# ============================================================================
+
+app = Server()
+
+@app.get("/")
+async def homepage():
+    html_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "index.html")
+    with open(html_path, "r", encoding="utf-8") as f:
+        return HTMLResponse(content=f.read())
+
+@app.get("/app_config")
+async def get_config():
+    """Return server configuration for frontend (e.g. LOW_VRAM mode)."""
+    return JSONResponse({"low_vram": LOW_VRAM})
+
+@app.get("/progress")
+async def progress_poll(request: Request):
+    """Polling endpoint for real-time progress updates during generation."""
+    session_id = request.query_params.get("session_id", "")
+    path = _progress_file(session_id)
+    try:
+        with open(path, 'r') as f:
+            data = json.load(f)
+        return JSONResponse(data)
+    except (FileNotFoundError, json.JSONDecodeError):
+        return JSONResponse({"stage": "Waiting...", "step": 0, "total": 0, "done": False})
+
+@app.api()
+@spaces.GPU(duration=30)
+def preprocess(image: FileData) -> FileData:
+    init_models()
+    img = Image.open(image["path"])
+    processed = pipeline.preprocess_image(img)
+    out_path = os.path.join(TMP_DIR, f"preprocessed_{int(time.time()*1000)}.png")
+    processed.save(out_path)
+    return FileData(path=out_path)
+
+@app.api()
+@spaces.GPU(duration=120)
+def generate_3d(
+    image: FileData, 
+    seed: int, 
+    resolution: int,
+    ss_guidance_strength: float = 7.5,
+    ss_guidance_rescale: float = 0.7,
+    ss_sampling_steps: int = 12,
+    ss_rescale_t: float = 5.0,
+    shape_slat_guidance_strength: float = 7.5,
+    shape_slat_guidance_rescale: float = 0.5,
+    shape_slat_sampling_steps: int = 12,
+    shape_slat_rescale_t: float = 3.0,
+    tex_slat_guidance_strength: float = 1.0,
+    tex_slat_guidance_rescale: float = 0.0,
+    tex_slat_sampling_steps: int = 12,
+    tex_slat_rescale_t: float = 3.0,
+    manual_fov: float = -1.0,
+    fov_unit: str = "deg",
+    session_id: str = "",
+) -> Dict:
+    init_models()
+    _reset_progress(session_id)
+    _update_progress("Preprocessing & Camera Estimation", 0, 1)
+    
+    torch.manual_seed(seed)
+    hr_resolution = int(resolution)
+    
+    img = Image.open(image["path"])
+    # Image is already preprocessed by /preprocess endpoint, use directly
+    image_preprocessed = img
+    temp_processed_path = os.path.join(TMP_DIR, f"temp_proc_{session_id[:8]}_{int(time.time()*1000)}.png")
+    image_preprocessed.save(temp_processed_path)
+    
+    if manual_fov > 0:
+        # Convert to radians based on unit
+        if fov_unit == "rad":
+            camera_angle_x = float(manual_fov)
+            fov_deg = math.degrees(manual_fov)
+        else:
+            camera_angle_x = math.radians(manual_fov)
+            fov_deg = float(manual_fov)
+        grid_point = torch.tensor([-1.0, 0.0, 0.0])
+        distance = distance_from_fov(
+            camera_angle_x, grid_point,
+            torch.tensor([0 - WILD_EXTEND_PIXEL, WILD_IMAGE_RESOLUTION - 1 + WILD_EXTEND_PIXEL]),
+            WILD_MESH_SCALE, WILD_IMAGE_RESOLUTION
+        )["distance_from_x"]
+        camera_params = {'camera_angle_x': camera_angle_x, 'distance': distance, 'mesh_scale': WILD_MESH_SCALE}
+        print(f"[Camera] Using manual FOV: {fov_deg:.2f}° ({camera_angle_x:.4f} rad), distance: {distance:.4f}")
+    else:
+        camera_params = get_camera_params_wild_moge(
+            temp_processed_path, device="cuda",
+            mesh_scale=WILD_MESH_SCALE, extend_pixel=WILD_EXTEND_PIXEL,
+            image_resolution=WILD_IMAGE_RESOLUTION,
+        )
+    _update_progress("Preprocessing & Camera Estimation", 1, 1)
+    
+    ss_sampler_override = {"steps": ss_sampling_steps, "guidance_strength": ss_guidance_strength,
+                           "guidance_rescale": ss_guidance_rescale, "rescale_t": ss_rescale_t}
+    shape_sampler_override = {"steps": shape_slat_sampling_steps, "guidance_strength": shape_slat_guidance_strength,
+                              "guidance_rescale": shape_slat_guidance_rescale, "rescale_t": shape_slat_rescale_t}
+    tex_sampler_override = {"steps": tex_slat_sampling_steps, "guidance_strength": tex_slat_guidance_strength,
+                            "guidance_rescale": tex_slat_guidance_rescale, "rescale_t": tex_slat_rescale_t}
+
+    pipeline_type = f"{hr_resolution}_cascade"
+    mesh_list, (shape_slat, tex_slat, res) = pipeline.run(
+        image_preprocessed,
+        camera_params=camera_params,
+        seed=seed,
+        sparse_structure_sampler_params=ss_sampler_override,
+        shape_slat_sampler_params=shape_sampler_override,
+        tex_slat_sampler_params=tex_sampler_override,
+        preprocess_image=False,
+        return_latent=True,
+        pipeline_type=pipeline_type,
+        max_num_tokens=CASCADE_MAX_NUM_TOKENS,
+    )
+    
+    mesh = mesh_list[0]
+    state_path = pack_state(shape_slat, tex_slat, res)
+    
+    _update_progress("Rendering views", 0, 1)
+    mesh.simplify(16777216)
+    cam_dist = camera_params['distance']
+    near = max(0.01, cam_dist - 2.0)
+    far = cam_dist + 10.0
+    if LOW_VRAM:
+        for v in envmap.values():
+            v.image = v.image.cuda()
+            if hasattr(v, '_nvdiffrec_envlight'):
+                del v._nvdiffrec_envlight
+    renders = render_utils.render_proj_aligned_video(
+        mesh, camera_angle_x=camera_params['camera_angle_x'],
+        distance=cam_dist, resolution=1024,
+        num_frames=STEPS, envmap=envmap,
+        near=near, far=far,
+    )
+    if LOW_VRAM:
+        for v in envmap.values():
+            if hasattr(v, '_nvdiffrec_envlight'):
+                del v._nvdiffrec_envlight
+            v.image = v.image.cpu()
+        torch.cuda.empty_cache()
+    _update_progress("Rendering views", 1, 1)
+    
+    # Save renders and return paths
+    render_files = {}
+    for mode_key, frames in renders.items():
+        mode_files = []
+        for i, frame in enumerate(frames):
+            p = os.path.abspath(os.path.join(TMP_DIR, f"render_{mode_key}_{i}_{int(time.time()*1000)}.jpg"))
+            Image.fromarray(frame).save(p, quality=85)
+            mode_files.append(FileData(path=p))
+        render_files[mode_key] = mode_files
+
+    _finish_progress()
+    return {
+        "render_paths": render_files,
+        "state_path": os.path.abspath(state_path),
+        "camera_angle_x": camera_params['camera_angle_x'],
+        "distance": camera_params['distance'],
+    }
+
+@app.api()
+@spaces.GPU(duration=240)
+def extract_glb_api(state_path: str, decimation_target: int, texture_size: int, session_id: str = "") -> FileData:
+    init_models()
+    _reset_progress(session_id)
+    _update_progress("Decoding latent", 0, 1)
+    
+    shape_slat, tex_slat, res = unpack_state(state_path)
+    mesh = pipeline.decode_latent(shape_slat, tex_slat, res)[0]
+    _update_progress("Decoding latent", 1, 1)
+    
+    glb = o_voxel.postprocess.to_glb(
+        vertices=mesh.vertices, faces=mesh.faces, attr_volume=mesh.attrs,
+        coords=mesh.coords, attr_layout=pipeline.pbr_attr_layout,
+        grid_size=res, aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
+        decimation_target=decimation_target, texture_size=texture_size,
+        remesh=True, remesh_band=1, remesh_project=0, use_tqdm=True,
+    )
+    rot = np.array([
+        [-1,  0,  0,  0],
+        [ 0,  0, -1,  0],
+        [ 0, -1,  0,  0],
+        [ 0,  0,  0,  1],
+    ], dtype=np.float64)
+    glb.apply_transform(rot)
+    
+    out_glb = os.path.join(TMP_DIR, f"result_{int(time.time()*1000)}.glb")
+    glb.export(out_glb, extension_webp=True)
+    _finish_progress()
+    return FileData(path=out_glb)
+
+# Mount assets and tmp for direct access
+app.mount("/assets", StaticFiles(directory="assets"), name="assets")
+app.mount("/tmp", StaticFiles(directory=TMP_DIR), name="tmp")
+
+if __name__ == "__main__":
+    import sys
+    parser = argparse.ArgumentParser(description="Pixal3D Demo Server")
+    parser.add_argument("--low_vram", action="store_true",
+                        help="Enable low-VRAM mode: models lazy-load to GPU per stage.")
+    args, remaining = parser.parse_known_args()
+    if args.low_vram:
+        LOW_VRAM = True
+
+    # Re-install utils3d as in original app.py
+    subprocess.run([
+        sys.executable, "-m", "pip", "install", "--force-reinstall", "--no-deps",
+        "https://github.com/LDYang694/Storages/releases/download/20260430/utils3d-0.0.2-py3-none-any.whl"
+    ], check=True)
+    
+    # Pre-initialize models before launching the server
+    init_models()
+    
+    app.launch(show_error=True, share=True)
\ No newline at end of file
diff --git a/assets/app/basecolor.png b/assets/app/basecolor.png
new file mode 100644
index 0000000..e7dbeaf
Binary files /dev/null and b/assets/app/basecolor.png differ
diff --git a/assets/app/clay.png b/assets/app/clay.png
new file mode 100644
index 0000000..e02866a
Binary files /dev/null and b/assets/app/clay.png differ
diff --git a/assets/app/hdri_city.png b/assets/app/hdri_city.png
new file mode 100644
index 0000000..43e1c2a
Binary files /dev/null and b/assets/app/hdri_city.png differ
diff --git a/assets/app/hdri_courtyard.png b/assets/app/hdri_courtyard.png
new file mode 100644
index 0000000..4261ad6
Binary files /dev/null and b/assets/app/hdri_courtyard.png differ
diff --git a/assets/app/hdri_forest.png b/assets/app/hdri_forest.png
new file mode 100644
index 0000000..7617fe1
Binary files /dev/null and b/assets/app/hdri_forest.png differ
diff --git a/assets/app/hdri_interior.png b/assets/app/hdri_interior.png
new file mode 100644
index 0000000..e00c1d6
Binary files /dev/null and b/assets/app/hdri_interior.png differ
diff --git a/assets/app/hdri_night.png b/assets/app/hdri_night.png
new file mode 100644
index 0000000..f0423d2
Binary files /dev/null and b/assets/app/hdri_night.png differ
diff --git a/assets/app/hdri_studio.png b/assets/app/hdri_studio.png
new file mode 100644
index 0000000..0f5a4e8
Binary files /dev/null and b/assets/app/hdri_studio.png differ
diff --git a/assets/app/hdri_sunrise.png b/assets/app/hdri_sunrise.png
new file mode 100644
index 0000000..9cee3bb
Binary files /dev/null and b/assets/app/hdri_sunrise.png differ
diff --git a/assets/app/hdri_sunset.png b/assets/app/hdri_sunset.png
new file mode 100644
index 0000000..bd67070
Binary files /dev/null and b/assets/app/hdri_sunset.png differ
diff --git a/assets/app/normal.png b/assets/app/normal.png
new file mode 100644
index 0000000..352e92b
Binary files /dev/null and b/assets/app/normal.png differ
diff --git a/assets/hdri/city.exr b/assets/hdri/city.exr
new file mode 100644
index 0000000..d922066
Binary files /dev/null and b/assets/hdri/city.exr differ
diff --git a/assets/hdri/courtyard.exr b/assets/hdri/courtyard.exr
new file mode 100644
index 0000000..b70a0e7
Binary files /dev/null and b/assets/hdri/courtyard.exr differ
diff --git a/assets/hdri/forest.exr b/assets/hdri/forest.exr
new file mode 100644
index 0000000..846b87d
Binary files /dev/null and b/assets/hdri/forest.exr differ
diff --git a/assets/hdri/interior.exr b/assets/hdri/interior.exr
new file mode 100644
index 0000000..92d40ca
Binary files /dev/null and b/assets/hdri/interior.exr differ
diff --git a/assets/hdri/license.txt b/assets/hdri/license.txt
new file mode 100644
index 0000000..eba5cb8
--- /dev/null
+++ b/assets/hdri/license.txt
@@ -0,0 +1,15 @@
+All HDRIs are licensed as CC0.
+
+These were created by Greg Zaal (Poly Haven https://polyhaven.com).
+Originals used for each HDRI:
+- City: https://polyhaven.com/a/portland_landing_pad
+- Courtyard: https://polyhaven.com/a/courtyard
+- Forest: https://polyhaven.com/a/ninomaru_teien
+- Interior: https://polyhaven.com/a/hotel_room
+- Night: Probably https://polyhaven.com/a/moonless_golf
+- Studio: Probably https://polyhaven.com/a/studio_small_01
+- Sunrise: https://polyhaven.com/a/spruit_sunrise
+- Sunset: https://polyhaven.com/a/venice_sunset
+
+1K resolution of each was taken, and compressed with oiiotool:
+oiiotool input.exr --ch R,G,B -d float --compression dwab:300 --clamp:min=0.0:max=32000.0 -o output.exr
diff --git a/assets/hdri/night.exr b/assets/hdri/night.exr
new file mode 100644
index 0000000..a207d99
Binary files /dev/null and b/assets/hdri/night.exr differ
diff --git a/assets/hdri/studio.exr b/assets/hdri/studio.exr
new file mode 100644
index 0000000..baf478d
Binary files /dev/null and b/assets/hdri/studio.exr differ
diff --git a/assets/hdri/sunrise.exr b/assets/hdri/sunrise.exr
new file mode 100644
index 0000000..985a9a5
Binary files /dev/null and b/assets/hdri/sunrise.exr differ
diff --git a/assets/hdri/sunset.exr b/assets/hdri/sunset.exr
new file mode 100644
index 0000000..e86206e
Binary files /dev/null and b/assets/hdri/sunset.exr differ
diff --git a/assets/images/0_img.png b/assets/images/0_img.png
new file mode 100644
index 0000000..99a344f
Binary files /dev/null and b/assets/images/0_img.png differ
diff --git a/assets/images/10_img.webp b/assets/images/10_img.webp
new file mode 100644
index 0000000..f268f16
Binary files /dev/null and b/assets/images/10_img.webp differ
diff --git a/assets/images/11_img.png b/assets/images/11_img.png
new file mode 100644
index 0000000..076b649
Binary files /dev/null and b/assets/images/11_img.png differ
diff --git a/assets/images/12_img.png b/assets/images/12_img.png
new file mode 100644
index 0000000..349308f
Binary files /dev/null and b/assets/images/12_img.png differ
diff --git a/assets/images/17_img.png b/assets/images/17_img.png
new file mode 100644
index 0000000..3f03f16
Binary files /dev/null and b/assets/images/17_img.png differ
diff --git a/assets/images/1_img.png b/assets/images/1_img.png
new file mode 100644
index 0000000..4fa969f
Binary files /dev/null and b/assets/images/1_img.png differ
diff --git a/assets/images/21_img.png b/assets/images/21_img.png
new file mode 100644
index 0000000..21fa5ec
Binary files /dev/null and b/assets/images/21_img.png differ
diff --git a/assets/images/3_img.webp b/assets/images/3_img.webp
new file mode 100644
index 0000000..b833f8e
Binary files /dev/null and b/assets/images/3_img.webp differ
diff --git a/assets/images/4_img.png b/assets/images/4_img.png
new file mode 100644
index 0000000..b0bd8de
Binary files /dev/null and b/assets/images/4_img.png differ
diff --git a/assets/images/5_img.webp b/assets/images/5_img.webp
new file mode 100644
index 0000000..0e5c91a
Binary files /dev/null and b/assets/images/5_img.webp differ
diff --git a/assets/images/6_img.png b/assets/images/6_img.png
new file mode 100644
index 0000000..2dd339b
Binary files /dev/null and b/assets/images/6_img.png differ
diff --git a/assets/images/7_img.png b/assets/images/7_img.png
new file mode 100644
index 0000000..fc61ebd
Binary files /dev/null and b/assets/images/7_img.png differ
diff --git a/assets/images/9_img.png b/assets/images/9_img.png
new file mode 100644
index 0000000..8b4e10b
Binary files /dev/null and b/assets/images/9_img.png differ
diff --git a/assets/images/s_13_img.jpg b/assets/images/s_13_img.jpg
new file mode 100644
index 0000000..921bf11
Binary files /dev/null and b/assets/images/s_13_img.jpg differ
diff --git a/assets/images/s_14_img.jpg b/assets/images/s_14_img.jpg
new file mode 100644
index 0000000..5ef414e
Binary files /dev/null and b/assets/images/s_14_img.jpg differ
diff --git a/assets/images/s_15_img.png b/assets/images/s_15_img.png
new file mode 100644
index 0000000..613072d
Binary files /dev/null and b/assets/images/s_15_img.png differ
diff --git a/assets/images/s_16_img.png b/assets/images/s_16_img.png
new file mode 100644
index 0000000..7b5a6fd
Binary files /dev/null and b/assets/images/s_16_img.png differ
diff --git a/assets/images/s_18_img.png b/assets/images/s_18_img.png
new file mode 100644
index 0000000..1aad6af
Binary files /dev/null and b/assets/images/s_18_img.png differ
diff --git a/assets/images/s_20_img.webp b/assets/images/s_20_img.webp
new file mode 100644
index 0000000..56d7fd5
Binary files /dev/null and b/assets/images/s_20_img.webp differ
diff --git a/assets/teaser.png b/assets/teaser.png
new file mode 100644
index 0000000..16ba234
Binary files /dev/null and b/assets/teaser.png differ
diff --git a/configs/gen/slat_flow_img2shape_dit_1_3B_256_bf16_proj_finetune.json b/configs/gen/slat_flow_img2shape_dit_1_3B_256_bf16_proj_finetune.json
new file mode 100644
index 0000000..0a75187
--- /dev/null
+++ b/configs/gen/slat_flow_img2shape_dit_1_3B_256_bf16_proj_finetune.json
@@ -0,0 +1,113 @@
+{
+    "_comment": "TRELLIS.2 SLat Shape Flow Model (256) with View-Aligned Projection - Finetune",
+    "models": {
+        "denoiser": {
+            "name": "ElasticSLatFlowModel",
+            "args": {
+                "resolution": 16,
+                "in_channels": 32,
+                "out_channels": 32,
+                "model_channels": 1536,
+                "cond_channels": 1024,
+                "num_blocks": 30,
+                "num_heads": 12,
+                "mlp_ratio": 5.3334,
+                "pe_mode": "rope",
+                "share_mod": true,
+                "initialization": "scaled",
+                "qk_rms_norm": true,
+                "qk_rms_norm_cross": true,
+                "image_attn_mode": "proj",
+                "proj_in_channels": 2048
+            }
+        }
+    },
+    "dataset": {
+        "name": "ViewImageConditionedSLatShapeView",
+        "args": {
+            "resolution": 256,
+            "image_size": 512,
+            "min_aesthetic_score": 4.5,
+            "max_tokens": 8192,
+            "num_views": 2,
+            "normalization": {
+                "mean": [
+                    0.781296, 0.018091, -0.495192, -0.558457, 1.060530, 0.093252, 1.518149, -0.933218,
+                    -0.732996, 2.604095, -0.118341, -2.143904, 0.495076, -2.179512, -2.130751, -0.996944,
+                    0.261421, -2.217463, 1.260067, -0.150213, 3.790713, 1.481266, -1.046058, -1.523667,
+                    -0.059621, 2.220780, 1.621212, 0.877230, 0.567247, -3.175944, -3.186688, 1.578665
+                ],
+                "std": [
+                    5.972266, 4.706852, 5.445010, 5.209927, 5.320220, 4.547237, 5.020802, 5.444004,
+                    5.226681, 5.683095, 4.831436, 5.286469, 5.652043, 5.367606, 5.525084, 4.730578,
+                    4.805265, 5.124013, 5.530808, 5.619001, 5.103930, 5.417670, 5.269677, 5.547194,
+                    5.634698, 5.235274, 6.110351, 5.511298, 6.237273, 4.879207, 5.347008, 5.405691
+                ]
+            },
+            "skip_aesthetic_score_datasets": ["texverse"],
+            "pretrained_slat_dec": "microsoft/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16"
+        }
+    },
+    "trainer": {
+        "name": "ImageConditionedProjSparseFlowMatchingCFGTrainer",
+        "args": {
+            "max_steps": 1000000,
+            "batch_size_per_gpu": 16,
+            "batch_split": 2,
+            "snapshot_batch_size": 1,
+            "snapshot_num_samples": 1,
+            "num_workers": 8,
+            "optimizer": {
+                "name": "AdamW",
+                "args": {
+                    "lr": 1e-4,
+                    "weight_decay": 0.01,
+                    "betas": [0.9, 0.95],
+                    "eps": 1e-8
+                }
+            },
+            "ema_rate": [
+                0.9999
+            ],
+            "mix_precision_mode": "amp",
+            "mix_precision_dtype": "bfloat16",
+            "elastic": {
+                "name": "LinearMemoryController",
+                "args": {
+                    "target_ratio": 0.75,
+                    "max_mem_ratio_start": 0.5
+                }
+            },
+            "grad_clip": {
+                "name": "AdaptiveGradClipper",
+                "args": {
+                    "max_norm": 1.0,
+                    "clip_percentile": 95
+                }
+            },
+           
+            "i_print": 500,
+            "i_log": 5,
+            "i_sample": 2000,
+            "i_save": 5000,
+            "p_uncond": 0.1,
+            "t_schedule": {
+                "name": "uniform",
+                "args": {}
+            },
+            "sigma_min": 1e-5,
+            "run_projection_test": false,
+            "image_cond_model": {
+                "name": "DinoV3ProjFeatureExtractor",
+                "args": {
+                    "model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
+                    "image_size": 512,
+                    "grid_resolution": 16,
+                    "use_naf_upsample": true,
+                    "naf_target_size": 256
+                },
+                "image_attn_mode": "proj"
+            }
+        }
+    }
+}
diff --git a/configs/gen/slat_flow_img2shape_dit_1_3B_256_bf16_proj_finetune_ft512.json b/configs/gen/slat_flow_img2shape_dit_1_3B_256_bf16_proj_finetune_ft512.json
new file mode 100644
index 0000000..34ff75b
--- /dev/null
+++ b/configs/gen/slat_flow_img2shape_dit_1_3B_256_bf16_proj_finetune_ft512.json
@@ -0,0 +1,115 @@
+{
+    "_comment": "TRELLIS.2 SLat Shape Flow Model (512) with View-Aligned Projection + NAF Upsample - Finetune from 256 Proj Checkpoint. NAF upsamples DINOv3 32x32 features using the input image as guide. proj_in_channels=2048 = concat(lr 1024, hr 1024).",
+    "models": {
+        "denoiser": {
+            "name": "ElasticSLatFlowModel",
+            "args": {
+                "resolution": 32,
+                "in_channels": 32,
+                "out_channels": 32,
+                "model_channels": 1536,
+                "cond_channels": 1024,
+                "num_blocks": 30,
+                "num_heads": 12,
+                "mlp_ratio": 5.3334,
+                "pe_mode": "rope",
+                "share_mod": true,
+                "initialization": "scaled",
+                "qk_rms_norm": true,
+                "qk_rms_norm_cross": true,
+                "image_attn_mode": "proj",
+                "proj_in_channels": 2048
+            }
+        }
+    },
+    "dataset": {
+        "name": "ViewImageConditionedSLatShapeView",
+        "args": {
+            "resolution": 512,
+            "image_size": 512,
+            "min_aesthetic_score": 4.5,
+            "max_tokens": 8192,
+            "num_views": 2,
+            "normalization": {
+                "mean": [
+                    0.781296, 0.018091, -0.495192, -0.558457, 1.060530, 0.093252, 1.518149, -0.933218,
+                    -0.732996, 2.604095, -0.118341, -2.143904, 0.495076, -2.179512, -2.130751, -0.996944,
+                    0.261421, -2.217463, 1.260067, -0.150213, 3.790713, 1.481266, -1.046058, -1.523667,
+                    -0.059621, 2.220780, 1.621212, 0.877230, 0.567247, -3.175944, -3.186688, 1.578665
+                ],
+                "std": [
+                    5.972266, 4.706852, 5.445010, 5.209927, 5.320220, 4.547237, 5.020802, 5.444004,
+                    5.226681, 5.683095, 4.831436, 5.286469, 5.652043, 5.367606, 5.525084, 4.730578,
+                    4.805265, 5.124013, 5.530808, 5.619001, 5.103930, 5.417670, 5.269677, 5.547194,
+                    5.634698, 5.235274, 6.110351, 5.511298, 6.237273, 4.879207, 5.347008, 5.405691
+                ]
+            },
+            "skip_aesthetic_score_datasets": ["texverse"],
+            "pretrained_slat_dec": "microsoft/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16"
+        }
+    },
+    "trainer": {
+        "name": "ImageConditionedProjSparseFlowMatchingCFGTrainer",
+        "args": {
+            "max_steps": 1000000,
+            "batch_size_per_gpu": 8,
+            "batch_split": 2,
+            "snapshot_batch_size": 1,
+            "snapshot_num_samples": 1,
+            "num_workers": 14,
+            "optimizer": {
+                "name": "AdamW",
+                "args": {
+                    "lr": 1e-4,
+                    "weight_decay": 0.01,
+                    "betas": [0.9, 0.95],
+                    "eps": 1e-8
+                }
+            },
+            "ema_rate": [
+                0.9999
+            ],
+            "mix_precision_mode": "amp",
+            "mix_precision_dtype": "bfloat16",
+            "elastic": {
+                "name": "LinearMemoryController",
+                "args": {
+                    "target_ratio": 0.75,
+                    "max_mem_ratio_start": 0.5
+                }
+            },
+            "grad_clip": {
+                "name": "AdaptiveGradClipper",
+                "args": {
+                    "max_norm": 1.0,
+                    "clip_percentile": 95
+                }
+            },
+            "finetune_ckpt": {
+                "denoiser": "<path_to_slat_img2shape_256_checkpoint>"
+            },
+            "i_print": 500,
+            "i_log": 5,
+            "i_sample": 1000,
+            "i_save": 5000,
+            "p_uncond": 0.1,
+            "t_schedule": {
+                "name": "uniform",
+                "args": {}
+            },
+            "sigma_min": 1e-5,
+            "run_projection_test": false,
+            "image_cond_model": {
+                "name": "DinoV3ProjFeatureExtractor",
+                "args": {
+                    "model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
+                    "image_size": 512,
+                    "grid_resolution": 32,
+                    "use_naf_upsample": true,
+                    "naf_target_size": 512
+                },
+                "image_attn_mode": "proj"
+            }
+        }
+    }
+}
diff --git a/configs/gen/slat_flow_img2shape_dit_1_3B_512_bf16_proj_finetune_ft1024.json b/configs/gen/slat_flow_img2shape_dit_1_3B_512_bf16_proj_finetune_ft1024.json
new file mode 100644
index 0000000..6458204
--- /dev/null
+++ b/configs/gen/slat_flow_img2shape_dit_1_3B_512_bf16_proj_finetune_ft1024.json
@@ -0,0 +1,115 @@
+{
+    "_comment": "TRELLIS.2 SLat Shape Flow Model (1024) with View-Aligned Projection + NAF Upsample - Finetune from 512 Proj Checkpoint. NAF upsamples DINOv3 32x32 features to higher resolution using the input image as guide. proj_in_channels=2048 = concat(lr 1024, hr 1024).",
+    "models": {
+        "denoiser": {
+            "name": "ElasticSLatFlowModel",
+            "args": {
+                "resolution": 64,
+                "in_channels": 32,
+                "out_channels": 32,
+                "model_channels": 1536,
+                "cond_channels": 1024,
+                "num_blocks": 30,
+                "num_heads": 12,
+                "mlp_ratio": 5.3334,
+                "pe_mode": "rope",
+                "share_mod": true,
+                "initialization": "scaled",
+                "qk_rms_norm": true,
+                "qk_rms_norm_cross": true,
+                "image_attn_mode": "proj",
+                "proj_in_channels": 2048
+            }
+        }
+    },
+    "dataset": {
+        "name": "ViewImageConditionedSLatShapeView",
+        "args": {
+            "resolution": 1024,
+            "image_size": 1024,
+            "min_aesthetic_score": 4.5,
+            "max_tokens": 32768,
+            "num_views": 2,
+            "normalization": {
+                "mean": [
+                    0.781296, 0.018091, -0.495192, -0.558457, 1.060530, 0.093252, 1.518149, -0.933218,
+                    -0.732996, 2.604095, -0.118341, -2.143904, 0.495076, -2.179512, -2.130751, -0.996944,
+                    0.261421, -2.217463, 1.260067, -0.150213, 3.790713, 1.481266, -1.046058, -1.523667,
+                    -0.059621, 2.220780, 1.621212, 0.877230, 0.567247, -3.175944, -3.186688, 1.578665
+                ],
+                "std": [
+                    5.972266, 4.706852, 5.445010, 5.209927, 5.320220, 4.547237, 5.020802, 5.444004,
+                    5.226681, 5.683095, 4.831436, 5.286469, 5.652043, 5.367606, 5.525084, 4.730578,
+                    4.805265, 5.124013, 5.530808, 5.619001, 5.103930, 5.417670, 5.269677, 5.547194,
+                    5.634698, 5.235274, 6.110351, 5.511298, 6.237273, 4.879207, 5.347008, 5.405691
+                ]
+            },
+            "skip_aesthetic_score_datasets": ["texverse"],
+            "pretrained_slat_dec": "microsoft/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16"
+        }
+    },
+    "trainer": {
+        "name": "ImageConditionedProjSparseFlowMatchingCFGTrainer",
+        "args": {
+            "max_steps": 1000000,
+            "batch_size_per_gpu": 2,
+            "batch_split": 2,
+            "snapshot_batch_size": 1,
+            "snapshot_num_samples": 1,
+            "num_workers": 14,
+            "optimizer": {
+                "name": "AdamW",
+                "args": {
+                    "lr": 2e-5,
+                    "weight_decay": 0.01,
+                    "betas": [0.9, 0.95],
+                    "eps": 1e-8
+                }
+            },
+            "ema_rate": [
+                0.9999
+            ],
+            "mix_precision_mode": "amp",
+            "mix_precision_dtype": "bfloat16",
+            "elastic": {
+                "name": "LinearMemoryController",
+                "args": {
+                    "target_ratio": 0.75,
+                    "max_mem_ratio_start": 0.25
+                }
+            },
+            "grad_clip": {
+                "name": "AdaptiveGradClipper",
+                "args": {
+                    "max_norm": 1.0,
+                    "clip_percentile": 95
+                }
+            },
+            "finetune_ckpt": {
+                "denoiser": "<path_to_slat_img2shape_512_checkpoint>"
+            },
+            "i_print": 500,
+            "i_log": 5,
+            "i_sample": 250,
+            "i_save": 1000,
+            "p_uncond": 0.1,
+            "t_schedule": {
+                "name": "uniform",
+                "args": {}
+            },
+            "sigma_min": 1e-5,
+            "run_projection_test": false,
+            "image_cond_model": {
+                "name": "DinoV3ProjFeatureExtractor",
+                "args": {
+                    "model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
+                    "image_size": 1024,
+                    "grid_resolution": 64,
+                    "use_naf_upsample": true,
+                    "naf_target_size": 512
+                },
+                "image_attn_mode": "proj"
+            }
+        }
+    }
+}
diff --git a/configs/gen/slat_flow_imgshape2tex_dit_1_3B_256_bf16_proj_finetune.json b/configs/gen/slat_flow_imgshape2tex_dit_1_3B_256_bf16_proj_finetune.json
new file mode 100644
index 0000000..3be6e89
--- /dev/null
+++ b/configs/gen/slat_flow_imgshape2tex_dit_1_3B_256_bf16_proj_finetune.json
@@ -0,0 +1,134 @@
+{
+    "_comment": "TRELLIS.2 SLat PBR/Texture Flow Model (256) with View-Aligned Projection + NAF Upsample - Finetune. NAF upsamples DINOv3 32x32 features to 128x128 using the input image as guide. proj_in_channels=2048 = concat(lr 1024, hr 1024).",
+    "models": {
+        "denoiser": {
+            "name": "ElasticSLatFlowModel",
+            "args": {
+                "resolution": 16,
+                "in_channels": 64,
+                "out_channels": 32,
+                "model_channels": 1536,
+                "cond_channels": 1024,
+                "num_blocks": 30,
+                "num_heads": 12,
+                "mlp_ratio": 5.3334,
+                "pe_mode": "rope",
+                "share_mod": true,
+                "initialization": "scaled",
+                "qk_rms_norm": true,
+                "qk_rms_norm_cross": true,
+                "image_attn_mode": "proj",
+                "proj_in_channels": 2048
+            }
+        }
+    },
+    "dataset": {
+        "name": "ViewImageConditionedSLatPbrView",
+        "args": {
+            "resolution": 256,
+            "image_size": 512,
+            "min_aesthetic_score": 4.5,
+            "max_tokens": 8192,
+            "num_views": 2,
+            "full_pbr": false,
+            "pbr_slat_normalization": {
+                "mean": [
+                    3.501659, 2.212398, 2.226094, 0.251093, -0.026248, -0.687364, 0.439898, -0.928075,
+                    0.029398, -0.339596, -0.869527, 1.038479, -0.972385, 0.126042, -1.129303, 0.455149,
+                    -1.209521, 2.069067, 0.544735, 2.569128, -0.323407, 2.293000, -1.925608, -1.217717,
+                    1.213905, 0.971588, -0.023631, 0.106750, 2.021786, 0.250524, -0.662387, -0.768862
+                ],
+                "std": [
+                    2.665652, 2.743913, 2.765121, 2.595319, 3.037293, 2.291316, 2.144656, 2.911822,
+                    2.969419, 2.501689, 2.154811, 3.163343, 2.621215, 2.381943, 3.186697, 3.021588,
+                    2.295916, 3.234985, 3.233086, 2.260140, 2.874801, 2.810596, 3.292720, 2.674999,
+                    2.680878, 2.372054, 2.451546, 2.353556, 2.995195, 2.379849, 2.786195, 2.775190
+                ]
+            },
+            "shape_slat_normalization": {
+                "mean": [
+                    0.781296, 0.018091, -0.495192, -0.558457, 1.060530, 0.093252, 1.518149, -0.933218,
+                    -0.732996, 2.604095, -0.118341, -2.143904, 0.495076, -2.179512, -2.130751, -0.996944,
+                    0.261421, -2.217463, 1.260067, -0.150213, 3.790713, 1.481266, -1.046058, -1.523667,
+                    -0.059621, 2.220780, 1.621212, 0.877230, 0.567247, -3.175944, -3.186688, 1.578665
+                ],
+                "std": [
+                    5.972266, 4.706852, 5.445010, 5.209927, 5.320220, 4.547237, 5.020802, 5.444004,
+                    5.226681, 5.683095, 4.831436, 5.286469, 5.652043, 5.367606, 5.525084, 4.730578,
+                    4.805265, 5.124013, 5.530808, 5.619001, 5.103930, 5.417670, 5.269677, 5.547194,
+                    5.634698, 5.235274, 6.110351, 5.511298, 6.237273, 4.879207, 5.347008, 5.405691
+                ]
+            },
+            "attrs": [
+                "base_color",
+                "metallic",
+                "roughness",
+                "alpha"
+            ],
+            "skip_aesthetic_score_datasets": ["texverse"],
+            "pretrained_pbr_slat_dec": "microsoft/TRELLIS.2-4B/ckpts/tex_dec_next_dc_f16c32_fp16",
+            "pretrained_shape_slat_dec": "microsoft/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16"
+        }
+    },
+    "trainer": {
+        "name": "ImageConditionedProjSparseFlowMatchingCFGTrainer",
+        "args": {
+            "max_steps": 1000000,
+            "batch_size_per_gpu": 16,
+            "batch_split": 1,
+            "snapshot_batch_size": 1,
+            "snapshot_num_samples": 1,
+            "num_workers": 16,
+            "optimizer": {
+                "name": "AdamW",
+                "args": {
+                    "lr": 1e-4,
+                    "weight_decay": 0.01,
+                    "betas": [0.9, 0.95],
+                    "eps": 1e-8
+                }
+            },
+            "ema_rate": [
+                0.9999
+            ],
+            "mix_precision_mode": "amp",
+            "mix_precision_dtype": "bfloat16",
+            "elastic": {
+                "name": "LinearMemoryController",
+                "args": {
+                    "target_ratio": 0.75,
+                    "max_mem_ratio_start": 0.5
+                }
+            },
+            "grad_clip": {
+                "name": "AdaptiveGradClipper",
+                "args": {
+                    "max_norm": 1.0,
+                    "clip_percentile": 95
+                }
+            },
+            "i_print": 500,
+            "i_log": 5,
+            "i_sample": 1000,
+            "i_save": 5000,
+            "p_uncond": 0.1,
+            "t_schedule": {
+                "name": "uniform",
+                "args": {}
+            },
+            "sigma_min": 1e-5,
+            "run_projection_test": false,
+            "image_cond_model": {
+                "name": "DinoV3ProjFeatureExtractor",
+                "args": {
+                    "model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
+                    "image_size": 512,
+                    "grid_resolution": 16,
+                    "use_naf_upsample": true,
+                    "naf_target_size": 256
+                },
+                "image_attn_mode": "proj"
+            }
+        }
+    }
+}
diff --git a/configs/gen/slat_flow_imgshape2tex_dit_1_3B_512_bf16_proj_finetune.json b/configs/gen/slat_flow_imgshape2tex_dit_1_3B_512_bf16_proj_finetune.json
new file mode 100644
index 0000000..0e3bd90
--- /dev/null
+++ b/configs/gen/slat_flow_imgshape2tex_dit_1_3B_512_bf16_proj_finetune.json
@@ -0,0 +1,137 @@
+{
+    "_comment": "TRELLIS.2 SLat PBR/Texture Flow Model (512) with View-Aligned Projection + NAF Upsample - Finetune from 256 Proj Checkpoint. NAF upsamples DINOv3 32x32 features to 256x256 using the input image as guide. proj_in_channels=2048 = concat(lr 1024, hr 1024).",
+    "models": {
+        "denoiser": {
+            "name": "ElasticSLatFlowModel",
+            "args": {
+                "resolution": 32,
+                "in_channels": 64,
+                "out_channels": 32,
+                "model_channels": 1536,
+                "cond_channels": 1024,
+                "num_blocks": 30,
+                "num_heads": 12,
+                "mlp_ratio": 5.3334,
+                "pe_mode": "rope",
+                "share_mod": true,
+                "initialization": "scaled",
+                "qk_rms_norm": true,
+                "qk_rms_norm_cross": true,
+                "image_attn_mode": "proj",
+                "proj_in_channels": 2048
+            }
+        }
+    },
+    "dataset": {
+        "name": "ViewImageConditionedSLatPbrView",
+        "args": {
+            "resolution": 512,
+            "image_size": 512,
+            "min_aesthetic_score": 4.5,
+            "max_tokens": 8192,
+            "num_views": 2,
+            "full_pbr": false,
+            "pbr_slat_normalization": {
+                "mean": [
+                    3.501659, 2.212398, 2.226094, 0.251093, -0.026248, -0.687364, 0.439898, -0.928075,
+                    0.029398, -0.339596, -0.869527, 1.038479, -0.972385, 0.126042, -1.129303, 0.455149,
+                    -1.209521, 2.069067, 0.544735, 2.569128, -0.323407, 2.293000, -1.925608, -1.217717,
+                    1.213905, 0.971588, -0.023631, 0.106750, 2.021786, 0.250524, -0.662387, -0.768862
+                ],
+                "std": [
+                    2.665652, 2.743913, 2.765121, 2.595319, 3.037293, 2.291316, 2.144656, 2.911822,
+                    2.969419, 2.501689, 2.154811, 3.163343, 2.621215, 2.381943, 3.186697, 3.021588,
+                    2.295916, 3.234985, 3.233086, 2.260140, 2.874801, 2.810596, 3.292720, 2.674999,
+                    2.680878, 2.372054, 2.451546, 2.353556, 2.995195, 2.379849, 2.786195, 2.775190
+                ]
+            },
+            "shape_slat_normalization": {
+                "mean": [
+                    0.781296, 0.018091, -0.495192, -0.558457, 1.060530, 0.093252, 1.518149, -0.933218,
+                    -0.732996, 2.604095, -0.118341, -2.143904, 0.495076, -2.179512, -2.130751, -0.996944,
+                    0.261421, -2.217463, 1.260067, -0.150213, 3.790713, 1.481266, -1.046058, -1.523667,
+                    -0.059621, 2.220780, 1.621212, 0.877230, 0.567247, -3.175944, -3.186688, 1.578665
+                ],
+                "std": [
+                    5.972266, 4.706852, 5.445010, 5.209927, 5.320220, 4.547237, 5.020802, 5.444004,
+                    5.226681, 5.683095, 4.831436, 5.286469, 5.652043, 5.367606, 5.525084, 4.730578,
+                    4.805265, 5.124013, 5.530808, 5.619001, 5.103930, 5.417670, 5.269677, 5.547194,
+                    5.634698, 5.235274, 6.110351, 5.511298, 6.237273, 4.879207, 5.347008, 5.405691
+                ]
+            },
+            "attrs": [
+                "base_color",
+                "metallic",
+                "roughness",
+                "alpha"
+            ],
+            "skip_aesthetic_score_datasets": ["texverse"],
+            "pretrained_pbr_slat_dec": "microsoft/TRELLIS.2-4B/ckpts/tex_dec_next_dc_f16c32_fp16",
+            "pretrained_shape_slat_dec": "microsoft/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16"
+        }
+    },
+    "trainer": {
+        "name": "ImageConditionedProjSparseFlowMatchingCFGTrainer",
+        "args": {
+            "max_steps": 1000000,
+            "batch_size_per_gpu": 8,
+            "batch_split": 2,
+            "snapshot_batch_size": 1,
+            "snapshot_num_samples": 1,
+            "num_workers": 14,
+            "optimizer": {
+                "name": "AdamW",
+                "args": {
+                    "lr": 1e-4,
+                    "weight_decay": 0.01,
+                    "betas": [0.9, 0.95],
+                    "eps": 1e-8
+                }
+            },
+            "ema_rate": [
+                0.9999
+            ],
+            "mix_precision_mode": "amp",
+            "mix_precision_dtype": "bfloat16",
+            "elastic": {
+                "name": "LinearMemoryController",
+                "args": {
+                    "target_ratio": 0.75,
+                    "max_mem_ratio_start": 0.5
+                }
+            },
+            "grad_clip": {
+                "name": "AdaptiveGradClipper",
+                "args": {
+                    "max_norm": 1.0,
+                    "clip_percentile": 95
+                }
+            },
+            "finetune_ckpt": {
+                "denoiser": "<path_to_slat_imgshape2tex_256_checkpoint>"
+            },
+            "i_print": 500,
+            "i_log": 5,
+            "i_sample": 1000,
+            "i_save": 5000,
+            "p_uncond": 0.1,
+            "t_schedule": {
+                "name": "uniform",
+                "args": {}
+            },
+            "sigma_min": 1e-5,
+            "run_projection_test": false,
+            "image_cond_model": {
+                "name": "DinoV3ProjFeatureExtractor",
+                "args": {
+                    "model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
+                    "image_size": 512,
+                    "grid_resolution": 32,
+                    "use_naf_upsample": true,
+                    "naf_target_size": 256
+                },
+                "image_attn_mode": "proj"
+            }
+        }
+    }
+}
diff --git a/configs/gen/slat_flow_imgshape2tex_dit_1_3B_512_bf16_proj_finetune_ft1024.json b/configs/gen/slat_flow_imgshape2tex_dit_1_3B_512_bf16_proj_finetune_ft1024.json
new file mode 100644
index 0000000..44701eb
--- /dev/null
+++ b/configs/gen/slat_flow_imgshape2tex_dit_1_3B_512_bf16_proj_finetune_ft1024.json
@@ -0,0 +1,137 @@
+{
+    "_comment": "TRELLIS.2 SLat PBR/Texture Flow Model (1024) with View-Aligned Projection + NAF Upsample - Finetune from 512 Proj Checkpoint. NAF upsamples DINOv3 32x32 features to 512x512 using the input image as guide. proj_in_channels=2048 = concat(lr 1024, hr 1024).",
+    "models": {
+        "denoiser": {
+            "name": "ElasticSLatFlowModel",
+            "args": {
+                "resolution": 64,
+                "in_channels": 64,
+                "out_channels": 32,
+                "model_channels": 1536,
+                "cond_channels": 1024,
+                "num_blocks": 30,
+                "num_heads": 12,
+                "mlp_ratio": 5.3334,
+                "pe_mode": "rope",
+                "share_mod": true,
+                "initialization": "scaled",
+                "qk_rms_norm": true,
+                "qk_rms_norm_cross": true,
+                "image_attn_mode": "proj",
+                "proj_in_channels": 2048
+            }
+        }
+    },
+    "dataset": {
+        "name": "ViewImageConditionedSLatPbrView",
+        "args": {
+            "resolution": 1024,
+            "image_size": 1024,
+            "min_aesthetic_score": 4.5,
+            "max_tokens": 32768,
+            "num_views": 2,
+            "full_pbr": true,
+            "pbr_slat_normalization": {
+                "mean": [
+                    3.501659, 2.212398, 2.226094, 0.251093, -0.026248, -0.687364, 0.439898, -0.928075,
+                    0.029398, -0.339596, -0.869527, 1.038479, -0.972385, 0.126042, -1.129303, 0.455149,
+                    -1.209521, 2.069067, 0.544735, 2.569128, -0.323407, 2.293000, -1.925608, -1.217717,
+                    1.213905, 0.971588, -0.023631, 0.106750, 2.021786, 0.250524, -0.662387, -0.768862
+                ],
+                "std": [
+                    2.665652, 2.743913, 2.765121, 2.595319, 3.037293, 2.291316, 2.144656, 2.911822,
+                    2.969419, 2.501689, 2.154811, 3.163343, 2.621215, 2.381943, 3.186697, 3.021588,
+                    2.295916, 3.234985, 3.233086, 2.260140, 2.874801, 2.810596, 3.292720, 2.674999,
+                    2.680878, 2.372054, 2.451546, 2.353556, 2.995195, 2.379849, 2.786195, 2.775190
+                ]
+            },
+            "shape_slat_normalization": {
+                "mean": [
+                    0.781296, 0.018091, -0.495192, -0.558457, 1.060530, 0.093252, 1.518149, -0.933218,
+                    -0.732996, 2.604095, -0.118341, -2.143904, 0.495076, -2.179512, -2.130751, -0.996944,
+                    0.261421, -2.217463, 1.260067, -0.150213, 3.790713, 1.481266, -1.046058, -1.523667,
+                    -0.059621, 2.220780, 1.621212, 0.877230, 0.567247, -3.175944, -3.186688, 1.578665
+                ],
+                "std": [
+                    5.972266, 4.706852, 5.445010, 5.209927, 5.320220, 4.547237, 5.020802, 5.444004,
+                    5.226681, 5.683095, 4.831436, 5.286469, 5.652043, 5.367606, 5.525084, 4.730578,
+                    4.805265, 5.124013, 5.530808, 5.619001, 5.103930, 5.417670, 5.269677, 5.547194,
+                    5.634698, 5.235274, 6.110351, 5.511298, 6.237273, 4.879207, 5.347008, 5.405691
+                ]
+            },
+            "attrs": [
+                "base_color",
+                "metallic",
+                "roughness",
+                "alpha"
+            ],
+            "skip_aesthetic_score_datasets": ["texverse"],
+            "pretrained_pbr_slat_dec": "microsoft/TRELLIS.2-4B/ckpts/tex_dec_next_dc_f16c32_fp16",
+            "pretrained_shape_slat_dec": "microsoft/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16"
+        }
+    },
+    "trainer": {
+        "name": "ImageConditionedProjSparseFlowMatchingCFGTrainer",
+        "args": {
+            "max_steps": 1000000,
+            "batch_size_per_gpu": 2,
+            "batch_split": 2,
+            "snapshot_batch_size": 1,
+            "snapshot_num_samples": 1,
+            "num_workers": 14,
+            "optimizer": {
+                "name": "AdamW",
+                "args": {
+                    "lr": 2e-5,
+                    "weight_decay": 0.01,
+                    "betas": [0.9, 0.95],
+                    "eps": 1e-8
+                }
+            },
+            "ema_rate": [
+                0.9999
+            ],
+            "mix_precision_mode": "amp",
+            "mix_precision_dtype": "bfloat16",
+            "elastic": {
+                "name": "LinearMemoryController",
+                "args": {
+                    "target_ratio": 0.75,
+                    "max_mem_ratio_start": 0.25
+                }
+            },
+            "grad_clip": {
+                "name": "AdaptiveGradClipper",
+                "args": {
+                    "max_norm": 1.0,
+                    "clip_percentile": 95
+                }
+            },
+            "finetune_ckpt": {
+                "denoiser": "<path_to_slat_imgshape2tex_512_checkpoint>"
+            },
+            "i_print": 500,
+            "i_log": 5,
+            "i_sample": -1,
+            "i_save": 250,
+            "p_uncond": 0.1,
+            "t_schedule": {
+                "name": "uniform",
+                "args": {}
+            },
+            "sigma_min": 1e-5,
+            "run_projection_test": false,
+            "image_cond_model": {
+                "name": "DinoV3ProjFeatureExtractor",
+                "args": {
+                    "model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
+                    "image_size": 1024,
+                    "grid_resolution": 64,
+                    "use_naf_upsample": true,
+                    "naf_target_size": 1024
+                },
+                "image_attn_mode": "proj"
+            }
+        }
+    }
+}
diff --git a/configs/gen/ss_flow_img_dit_1_3B_32_bf16_proj_finetune.json b/configs/gen/ss_flow_img_dit_1_3B_32_bf16_proj_finetune.json
new file mode 100644
index 0000000..2fd59dc
--- /dev/null
+++ b/configs/gen/ss_flow_img_dit_1_3B_32_bf16_proj_finetune.json
@@ -0,0 +1,91 @@
+{
+    "_comment": "TRELLIS.2 Sparse Structure Flow Model with View-Aligned Projection Conditioning (32 resolution, view-based)",
+    "models": {
+        "denoiser": {
+            "name": "SparseStructureFlowModel",
+            "args": {
+                "resolution": 8,
+                "in_channels": 8,
+                "out_channels": 8,
+                "model_channels": 1536,
+                "cond_channels": 1024,
+                "num_blocks": 30,
+                "num_heads": 12,
+                "mlp_ratio": 5.3334,
+                "pe_mode": "rope",
+                "share_mod": true,
+                "initialization": "scaled",
+                "qk_rms_norm": true,
+                "qk_rms_norm_cross": true,
+                "image_attn_mode": "proj"
+            }
+        }
+    },
+    "dataset": {
+        "name": "ViewImageConditionedSparseStructureLatentView",
+        "args": {
+            "min_aesthetic_score": 4.5,
+            "image_size": 512,
+            "num_views": 2,
+            "load_camera_info": true,
+            "skip_aesthetic_score_datasets": ["texverse"],
+            "pretrained_ss_dec": "microsoft/TRELLIS-image-large/ckpts/ss_dec_conv3d_16l8_fp16"
+        }
+    },
+    "trainer": {
+        "name": "ImageConditionedProjFlowMatchingCFGTrainer",
+        "args": {
+            "debug": false,
+            "max_steps": 1000000,
+            "batch_size_per_gpu": 16,
+            "batch_split": 1,
+            "snapshot_batch_size": 1,
+            "snapshot_num_samples": 1,
+            "num_workers": 14,
+            "optimizer": {
+                "name": "AdamW",
+                "args": {
+                    "lr": 1e-4,
+                    "weight_decay": 0.01,
+                    "betas": [0.9, 0.95],
+                    "eps": 1e-8
+                }
+            },
+            "ema_rate": [
+                0.9999
+            ],
+            "mix_precision_mode": "amp",
+            "mix_precision_dtype": "bfloat16",
+            "grad_clip": {
+                "name": "AdaptiveGradClipper",
+                "args": {
+                    "max_norm": 1.0,
+                    "clip_percentile": 95
+                }
+            },
+            "i_print": 500,
+            "i_log": 5,
+            "i_sample": 500,
+            "i_save": 5000,
+            "p_uncond": 0.1,
+            "t_schedule": {
+                "name": "logitNormal",
+                "args": {
+                    "mean": 1.0,
+                    "std": 1.0
+                }
+            },
+            "sigma_min": 1e-5,
+            "run_projection_test": false,
+            "image_cond_model": {
+                "name": "DinoV3ProjFeatureExtractor",
+                "args": {
+                    "model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
+                    "image_size": 512,
+                    "grid_resolution": 8
+                },
+                "image_attn_mode": "proj"
+            }
+        }
+    }
+}
diff --git a/configs/gen/ss_flow_img_dit_1_3B_32_bf16_proj_finetune_ft64.json b/configs/gen/ss_flow_img_dit_1_3B_32_bf16_proj_finetune_ft64.json
new file mode 100644
index 0000000..ac5bcc8
--- /dev/null
+++ b/configs/gen/ss_flow_img_dit_1_3B_32_bf16_proj_finetune_ft64.json
@@ -0,0 +1,94 @@
+{
+    "_comment": "TRELLIS.2 Sparse Structure Flow Model (64) with View-Aligned Projection - Finetune from 32 Proj Checkpoint",
+    "models": {
+        "denoiser": {
+            "name": "SparseStructureFlowModel",
+            "args": {
+                "resolution": 16,
+                "in_channels": 8,
+                "out_channels": 8,
+                "model_channels": 1536,
+                "cond_channels": 1024,
+                "num_blocks": 30,
+                "num_heads": 12,
+                "mlp_ratio": 5.3334,
+                "pe_mode": "rope",
+                "share_mod": true,
+                "initialization": "scaled",
+                "qk_rms_norm": true,
+                "qk_rms_norm_cross": true,
+                "image_attn_mode": "proj"
+            }
+        }
+    },
+    "dataset": {
+        "name": "ViewImageConditionedSparseStructureLatentView",
+        "args": {
+            "min_aesthetic_score": 4.5,
+            "image_size": 512,
+            "num_views": 2,
+            "load_camera_info": true,
+            "skip_aesthetic_score_datasets": ["texverse"],
+            "pretrained_ss_dec": "microsoft/TRELLIS-image-large/ckpts/ss_dec_conv3d_16l8_fp16"
+        }
+    },
+    "trainer": {
+        "name": "ImageConditionedProjFlowMatchingCFGTrainer",
+        "args": {
+            "debug": false,
+            "max_steps": 1000000,
+            "batch_size_per_gpu": 8,
+            "batch_split": 2,
+            "snapshot_batch_size": 1,
+            "snapshot_num_samples": 1,
+            "num_workers": 14,
+            "optimizer": {
+                "name": "AdamW",
+                "args": {
+                    "lr": 1e-4,
+                    "weight_decay": 0.01,
+                    "betas": [0.9, 0.95],
+                    "eps": 1e-8
+                }
+            },
+            "ema_rate": [
+                0.9999
+            ],
+            "mix_precision_mode": "amp",
+            "mix_precision_dtype": "bfloat16",
+            "grad_clip": {
+                "name": "AdaptiveGradClipper",
+                "args": {
+                    "max_norm": 1.0,
+                    "clip_percentile": 95
+                }
+            },
+            "finetune_ckpt": {
+                "denoiser": "<path_to_ss_flow_32_checkpoint>"
+            },
+            "i_print": 500,
+            "i_log": 5,
+            "i_sample": 2000,
+            "i_save": 5000,
+            "p_uncond": 0.1,
+            "t_schedule": {
+                "name": "logitNormal",
+                "args": {
+                    "mean": 1.0,
+                    "std": 1.0
+                }
+            },
+            "sigma_min": 1e-5,
+            "run_projection_test": false,
+            "image_cond_model": {
+                "name": "DinoV3ProjFeatureExtractor",
+                "args": {
+                    "model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
+                    "image_size": 512,
+                    "grid_resolution": 16
+                },
+                "image_attn_mode": "proj"
+            }
+        }
+    }
+}
diff --git a/data_toolkit/README.md b/data_toolkit/README.md
new file mode 100644
index 0000000..b95ff80
--- /dev/null
+++ b/data_toolkit/README.md
@@ -0,0 +1,212 @@
+# Dataset Preparation Toolkit
+
+This toolkit provides a comprehensive pipeline for preparing 3D datasets, including downloading, processing, voxelizing, and latent encoding for SC-VAE and Flow Model training.
+
+This toolkit is built upon and extended from the data processing scripts of [TRELLIS.2](https://github.com/microsoft/TRELLIS2). We gratefully acknowledge the TRELLIS.2 team for open-sourcing their data preparation pipeline, which served as the foundation for this work. Our extensions include view-aligned voxelization and latent encoding for Pixal3D.
+
+### Step 1: Install Dependencies
+
+Initialize the environment and install necessary dependencies:
+
+```bash
+. ./data_toolkit/setup.sh
+```
+
+### Step 2: Initialize Metadata
+
+Before processing, load the dataset metadata.
+
+```bash
+python data_toolkit/build_metadata.py <SUBSET> --root <ROOT> [--source <SOURCE>]
+```
+
+**Arguments:**
+- `SUBSET`: Target dataset subset. Options: `ObjaverseXL`, `ABO`, `HSSD`, `TexVerse` (Training sets); `SketchfabPicked`, `Toys4k` (Test sets).
+- `ROOT`: Root directory to save the data.
+- `SOURCE`: Data source (Required if `SUBSET` is `ObjaverseXL`). Options: `sketchfab`, `github`.
+
+**Example:**
+Load metadata for `ObjaverseXL` (sketchfab) and save to `datasets/ObjaverseXL_sketchfab`:
+```bash
+python data_toolkit/build_metadata.py ObjaverseXL --source sketchfab --root datasets/ObjaverseXL_sketchfab
+```
+
+### Step 3: Download Data
+
+Download the 3D assets to the local storage.
+
+```bash
+python data_toolkit/download.py <SUBSET> --root <ROOT> [--rank <RANK> --world_size <WORLD_SIZE>]
+```
+
+**Arguments:**
+- `RANK` / `WORLD_SIZE`: Parameters for multi-node distributed downloading.
+
+**Example:**
+To download the `ObjaverseXL` subset:
+
+> **Note:** The example below sets a large `WORLD_SIZE` (160,000) for demonstration purposes, meaning only a tiny fraction of the dataset will be downloaded by this single process.
+
+```bash
+python data_toolkit/download.py ObjaverseXL --root datasets/ObjaverseXL_sketchfab --world_size 160000
+```
+
+*Attention: Some datasets may require an interactive Hugging Face login or manual steps. Please follow any on-screen instructions.*
+
+**Update Metadata:**
+After downloading, update the metadata registry:
+```bash
+python data_toolkit/build_metadata.py ObjaverseXL --root datasets/ObjaverseXL_sketchfab
+```
+
+If download records are missing but files already exist locally, use `--from_file` to scan and rebuild:
+```bash
+python data_toolkit/build_metadata.py ObjaverseXL --root datasets/ObjaverseXL_sketchfab --from_file
+```
+
+### Step 4: Process Mesh and PBR Textures
+
+Standardize 3D assets by dumping mesh and PBR textures.
+*Note: This process utilizes the CPU.*
+
+```bash
+# Dump Meshes
+python data_toolkit/dump_mesh.py <SUBSET> --root <ROOT> [--rank <RANK> --world_size <WORLD_SIZE>]
+
+# Dump PBR Textures
+python data_toolkit/dump_pbr.py <SUBSET> --root <ROOT> [--rank <RANK> --world_size <WORLD_SIZE>]
+
+# Get statistics of the asset
+python data_toolkit/asset_stats.py --root <ROOT> [--rank <RANK> --world_size <WORLD_SIZE>]
+```
+
+**Example:**
+```bash
+python data_toolkit/dump_mesh.py ObjaverseXL --root datasets/ObjaverseXL_sketchfab
+python data_toolkit/dump_pbr.py ObjaverseXL --root datasets/ObjaverseXL_sketchfab
+python data_toolkit/asset_stats.py --root datasets/ObjaverseXL_sketchfab
+```
+
+**Update Metadata:**
+```bash
+python data_toolkit/build_metadata.py ObjaverseXL --root datasets/ObjaverseXL_sketchfab
+```
+
+### Step 5: Render Image Conditions
+
+Render multi-view images for each asset. These are used both as image conditions for the generator and as camera transforms for view-aligned processing in subsequent steps.
+*Note: Blender and Pillow will be automatically installed on first run.*
+
+```bash
+python data_toolkit/render_cond.py <SUBSET> --root <ROOT> [--num_views <NUM_VIEWS>] [--rank <RANK> --world_size <WORLD_SIZE>]
+```
+
+**Arguments:**
+- `NUM_VIEWS`: Number of views to render per asset. Default is `2`.
+
+**Example:**
+```bash
+python data_toolkit/render_cond.py ObjaverseXL --root datasets/ObjaverseXL_sketchfab
+```
+
+**Update Metadata:**
+```bash
+python data_toolkit/build_metadata.py ObjaverseXL --root datasets/ObjaverseXL_sketchfab
+```
+
+### Step 6: Convert to View-Aligned O-Voxels
+
+Convert the processed meshes and textures into view-aligned O-Voxels format. Each asset is transformed according to camera views from Step 5, producing per-view voxel representations.
+*Note: This process utilizes the CPU.*
+
+```bash
+python data_toolkit/dual_grid_view.py <SUBSET> --root <ROOT> [--rank <RANK> --world_size <WORLD_SIZE>] [--resolution <RESOLUTION>] [--view_indices <VIEW_INDICES>]
+
+python data_toolkit/voxelize_pbr_view.py <SUBSET> --root <ROOT> [--rank <RANK> --world_size <WORLD_SIZE>] [--resolution <RESOLUTION>] [--view_indices <VIEW_INDICES>]
+```
+
+**Arguments:**
+- `RESOLUTION`: Target resolutions for O-Voxels, comma-separated (e.g., `256,512,1024`). Default is `256`.
+- `VIEW_INDICES`: Specific view indices to process (e.g., `0,1,2` or `0-5`). Default processes all available views.
+
+**Example:**
+Convert `ObjaverseXL` to resolution 256 for views 0-1:
+```bash
+python data_toolkit/dual_grid_view.py ObjaverseXL --root datasets/ObjaverseXL_sketchfab --resolution 256 --view_indices 0-1
+python data_toolkit/voxelize_pbr_view.py ObjaverseXL --root datasets/ObjaverseXL_sketchfab --resolution 256 --view_indices 0-1
+```
+
+**Update Metadata:**
+```bash
+python data_toolkit/build_metadata.py ObjaverseXL --root datasets/ObjaverseXL_sketchfab
+```
+
+### At this point, the dataset is ready for SC-VAE Training
+
+### Step 7: Encode View-Aligned Latents
+
+Encode view-aligned sparse structures into latents to train the first-stage generator. Each step produces per-view latent files.
+
+```bash
+# 1. Encode Shape Latents (multi-view)
+python data_toolkit/encode_shape_latent_view.py --root <ROOT> [--rank <RANK> --world_size <WORLD_SIZE>] [--resolution <RESOLUTION>] [--view_indices <VIEW_INDICES>]
+
+# 2. Encode PBR Latents (view-aligned)
+python data_toolkit/encode_pbr_latent_view.py --root <ROOT> [--rank <RANK> --world_size <WORLD_SIZE>] [--resolution <RESOLUTION>] [--view_indices <VIEW_INDICES>]
+
+# 3. Update Metadata (Required before next step)
+python data_toolkit/build_metadata.py <SUBSET> --root <ROOT>
+
+# 4. Encode Sparse Structure (SS) Latents (multi-view)
+python data_toolkit/encode_ss_latent_view.py --root <ROOT> --shape_latent_name <SHAPE_LATENT_NAME> [--rank <RANK> --world_size <WORLD_SIZE>] [--resolution <SS_RESOLUTION>] [--view_indices <VIEW_INDICES>]
+```
+
+**Arguments:**
+- `RESOLUTION`: Input O-Voxel resolution. Default is `1024`.
+- `SS_RESOLUTION`: Resolution for sparse structures. Default is `64`.
+- `SHAPE_LATENT_NAME`: The specific version name of the shape latent (use the `_view` variant name).
+- `VIEW_INDICES`: Specific view indices to process (e.g., `0,1,2` or `0-5`).
+
+**Example:**
+```bash
+python data_toolkit/encode_shape_latent_view.py --root datasets/ObjaverseXL_sketchfab --resolution 512 --view_indices 0-1
+python data_toolkit/encode_pbr_latent_view.py --root datasets/ObjaverseXL_sketchfab --resolution 512 --view_indices 0-1
+python data_toolkit/encode_shape_latent_view.py --root datasets/ObjaverseXL_sketchfab --resolution 1024 --view_indices 0-1
+python data_toolkit/encode_pbr_latent_view.py --root datasets/ObjaverseXL_sketchfab --resolution 1024 --view_indices 0-1
+
+# Update metadata
+python data_toolkit/build_metadata.py ObjaverseXL --root datasets/ObjaverseXL_sketchfab
+
+# Encode SS Latents (view-aligned)
+python data_toolkit/encode_ss_latent_view.py --root datasets/ObjaverseXL_sketchfab --shape_latent_name shape_enc_next_dc_f16c32_fp16_1024_view --resolution 64 --view_indices 0-1
+
+# Final Metadata Update
+python data_toolkit/build_metadata.py ObjaverseXL --root datasets/ObjaverseXL_sketchfab
+```
+
+### Step 8: Visualize Decoded Latents (Optional)
+
+Decode latent files back to meshes, export GLB, and render a front-view image for visual inspection.
+
+**Shape Latent Visualization:**
+```bash
+python data_toolkit/visualize_shape_latent.py \
+    --root datasets/ObjaverseXL_sketchfab \
+    --sha256 <SHA256_HASH> \
+    --resolution 1024 \
+    --view_idx 0
+```
+
+**PBR Latent Visualization (shape + texture):**
+```bash
+python data_toolkit/visualize_pbr_latent.py \
+    --root datasets/ObjaverseXL_sketchfab \
+    --sha256 <SHA256_HASH> \
+    --resolution 1024 \
+    --view_idx 0
+```
+
+Outputs are saved to `<ROOT>/vis/<SHA256>/` (shape) or `<ROOT>/vis_pbr/<SHA256>/` (PBR), including:
+- Decoded GLB mesh (with PBR textures for PBR variant)
+- Front-view rendered images (normal/depth for shape; shaded/base_color/normal etc. for PBR)
+- Copied condition renders and camera transforms from Step 5
diff --git a/data_toolkit/asset_stats.py b/data_toolkit/asset_stats.py
new file mode 100644
index 0000000..8e26933
--- /dev/null
+++ b/data_toolkit/asset_stats.py
@@ -0,0 +1,132 @@
+import os
+import argparse
+import pickle
+from tqdm import tqdm
+import pandas as pd
+from easydict import EasyDict as edict
+from concurrent.futures import ThreadPoolExecutor
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--root', type=str, required=True,
+                        help='Directory to save the metadata')
+    parser.add_argument('--mesh_dump_root', type=str, default=None,
+                        help='Directory to save the mesh dumps')
+    parser.add_argument('--pbr_dump_root', type=str, default=None,
+                        help='Directory to save the pbr dumps')
+    parser.add_argument('--instances', type=str, default=None,
+                        help='Instances to process')
+    parser.add_argument('--rank', type=int, default=0)
+    parser.add_argument('--world_size', type=int, default=1)
+    parser.add_argument('--max_workers', type=int, default=0)
+    opt = parser.parse_args()
+    opt = edict(vars(opt))
+    opt.mesh_dump_root = opt.mesh_dump_root or opt.root
+    opt.pbr_dump_root = opt.pbr_dump_root or opt.root
+
+    os.makedirs(os.path.join(opt.root, 'asset_stats', 'new_records'), exist_ok=True)
+
+    # get file list
+    if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
+        raise ValueError('metadata.csv not found')
+    metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
+    if os.path.exists(os.path.join(opt.root, 'asset_stats', 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.root, 'asset_stats','metadata.csv')).set_index('sha256'))
+    if os.path.exists(os.path.join(opt.mesh_dump_root, 'mesh_dumps', 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.mesh_dump_root, 'mesh_dumps','metadata.csv')).set_index('sha256'))
+    if os.path.exists(os.path.join(opt.pbr_dump_root, 'pbr_dumps', 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.pbr_dump_root, 'pbr_dumps', 'metadata.csv')).set_index('sha256'))
+    metadata = metadata.reset_index()
+    if opt.instances is None:
+        if 'num_faces' in metadata.columns:
+            metadata = metadata[metadata['num_faces'].isnull()]
+        metadata = metadata[(metadata['mesh_dumped'] == True) | (metadata['pbr_dumped'] == True)]
+    else:
+        if os.path.exists(opt.instances):
+            with open(opt.instances, 'r') as f:
+                instances = f.read().splitlines()
+        else:
+            instances = opt.instances.split(',')
+        metadata = metadata[metadata['sha256'].isin(instances)]
+
+    start = len(metadata) * opt.rank // opt.world_size
+    end = len(metadata) * (opt.rank + 1) // opt.world_size
+    metadata = metadata[start:end]
+       
+    print(f'Processing {len(metadata)} objects...')
+
+    # process objects
+    records = []
+    with ThreadPoolExecutor(max_workers=opt.max_workers or os.cpu_count()) as executor, \
+         tqdm(total=len(metadata), desc='Processing objects') as pbar:
+        def worker(metadatum):
+            try:
+                sha256 = metadatum['sha256']
+                if metadatum['pbr_dumped'] == True:
+                    with open(os.path.join(opt.pbr_dump_root, 'pbr_dumps', f'{sha256}.pickle'), 'rb') as f:
+                        dump = pickle.load(f)
+
+                        num_faces = 0
+                        num_vertices = 0
+                        for obj in dump['objects']:
+                            if obj['vertices'].size == 0 or obj['faces'].size == 0:
+                                continue
+                            num_faces += obj['faces'].shape[0]
+                            num_vertices += obj['vertices'].shape[0]
+
+                        num_basecolor_tex = 0
+                        num_metallic_tex = 0
+                        num_roughness_tex = 0
+                        num_alpha_tex = 0
+                        for mat in dump['materials']:
+                            if mat['baseColorTexture'] is not None:
+                                num_basecolor_tex += 1
+                            if mat['metallicTexture'] is not None:
+                                num_metallic_tex += 1
+                            if mat['roughnessTexture'] is not None:
+                                num_roughness_tex += 1
+                            if mat['alphaTexture'] is not None:
+                                num_alpha_tex += 1
+
+                        record = {
+                            'sha256': sha256,
+                            'num_faces': num_faces,
+                            'num_vertices': num_vertices,
+                            'num_basecolor_tex': num_basecolor_tex,
+                            'num_metallic_tex': num_metallic_tex,
+                            'num_roughness_tex': num_roughness_tex,
+                            'num_alpha_tex': num_alpha_tex,
+                        }
+                        records.append(record)
+                else:
+                    with open(os.path.join(opt.mesh_dump_root,'mesh_dumps', f'{sha256}.pickle'), 'rb') as f:
+                        dump = pickle.load(f)
+
+                        num_faces = 0
+                        num_vertices = 0
+                        for obj in dump['objects']:
+                            if obj['vertices'].size == 0 or obj['faces'].size == 0:
+                                continue
+                            num_faces += obj['faces'].shape[0]
+                            num_vertices += obj['vertices'].shape[0]
+
+                        record = {
+                            'sha256': sha256,
+                            'num_faces': num_faces,
+                            'num_vertices': num_vertices,
+                        }
+                        records.append(record)
+                pbar.update()
+            except Exception as e:
+                print(f'Error processing {sha256}: {e}')
+                pbar.update()                
+
+        for metadatum in metadata.to_dict('records'):
+            executor.submit(worker, metadatum)
+
+        executor.shutdown(wait=True)
+
+    # save records
+    records = pd.DataFrame.from_records(records)
+    records.to_csv(os.path.join(opt.root, 'asset_stats', 'new_records', f'part_{opt.rank}.csv'), index=False)
diff --git a/data_toolkit/blender_script/dump_mesh.py b/data_toolkit/blender_script/dump_mesh.py
new file mode 100644
index 0000000..eaa5d55
--- /dev/null
+++ b/data_toolkit/blender_script/dump_mesh.py
@@ -0,0 +1,242 @@
+import argparse, sys, os, math, io
+from typing import *
+import bpy
+import bmesh
+from mathutils import Vector, Matrix
+import numpy as np
+import pickle
+
+
+"""=============== BLENDER ==============="""
+
+IMPORT_FUNCTIONS: Dict[str, Callable] = {
+    "obj": bpy.ops.import_scene.obj if bpy.app.version[0] < 4 else bpy.ops.wm.obj_import,
+    "glb": bpy.ops.import_scene.gltf,
+    "gltf": bpy.ops.import_scene.gltf,
+    "usd": bpy.ops.import_scene.usd,
+    "fbx": bpy.ops.import_scene.fbx,
+    "stl": bpy.ops.import_mesh.stl if bpy.app.version[0] < 4 else bpy.ops.wm.stl_import,
+    "usda": bpy.ops.import_scene.usda,
+    "dae": bpy.ops.wm.collada_import,
+    "ply": bpy.ops.import_mesh.ply if bpy.app.version[0] < 4 else bpy.ops.wm.ply_import,
+    "abc": bpy.ops.wm.alembic_import,
+    "blend": bpy.ops.wm.append,
+}
+
+
+def init_scene() -> None:
+    """Resets the scene to a clean state.
+
+    Returns:
+        None
+    """
+    # delete everything
+    for obj in bpy.data.objects:
+        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 = object_path.split(".")[-1].lower()
+    if file_extension is None:
+        raise ValueError(f"Unsupported file type: {object_path}")
+
+    if file_extension == "usdz":
+        # install usdz io package
+        dirname = os.path.dirname(os.path.realpath(__file__))
+        usdz_package = os.path.join(dirname, "io_scene_usdz.zip")
+        bpy.ops.preferences.addon_install(filepath=usdz_package)
+        # enable it
+        addon_name = "io_scene_usdz"
+        bpy.ops.preferences.addon_enable(module=addon_name)
+        # import the usdz
+        from io_scene_usdz.import_usdz import import_usdz
+
+        import_usdz(context, filepath=object_path, materials=True, animations=True)
+        return None
+
+    # load from existing import functions
+    import_function = IMPORT_FUNCTIONS[file_extension]
+
+    print(f"Loading object from {object_path}")
+    if file_extension == "blend":
+        import_function(directory=object_path, link=False)
+    elif file_extension in {"glb", "gltf"}:
+        import_function(filepath=object_path, merge_vertices=True, import_shading='NORMALS', bone_heuristic='TEMPERANCE')
+    else:
+        import_function(filepath=object_path)
+        
+        
+def delete_invisible_objects() -> None:
+    """Deletes all invisible objects in the scene.
+
+    Returns:
+        None
+    """
+    to_remove = []
+    for obj in bpy.context.scene.objects:
+        if obj.hide_viewport or obj.hide_render:
+            obj.hide_viewport = False
+            obj.hide_render = False
+            obj.hide_select = False
+            to_remove.append(obj)
+    for obj in to_remove:
+        bpy.data.objects.remove(obj, do_unlink=True)
+
+    # Delete invisible collections
+    invisible_collections = [col for col in bpy.data.collections if col.hide_viewport]
+    for col in invisible_collections:
+        bpy.data.collections.remove(col)
+      
+
+def scene_bbox() -> 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)
+
+    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
+    scene_meshes = [obj for obj in bpy.context.scene.objects.values() if isinstance(obj.data, bpy.types.Mesh)]
+    for obj in scene_meshes:
+        found = True
+        for coord in obj.bound_box:
+            coord = Vector(coord)
+            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 normalize_scene() -> Tuple[float, Vector]:
+    """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:
+        Tuple[float, Vector]: The scale factor and the offset applied to the scene.
+    """
+    scene_root_objects = [obj for obj in bpy.context.scene.objects.values() if not obj.parent]
+    if len(scene_root_objects) > 1:
+        # create an empty object to be used as a parent for all root objects
+        scene = bpy.data.objects.new("ParentEmpty", None)
+        bpy.context.scene.collection.objects.link(scene)
+
+        # parent all root objects to the empty object
+        for obj in scene_root_objects:
+            obj.parent = scene
+    else:
+        scene = scene_root_objects[0]
+
+    bbox_min, bbox_max = scene_bbox()
+    scale = 1 / max(bbox_max - bbox_min)
+    scene.scale = scene.scale * scale
+
+    # Apply scale to matrix_world.
+    bpy.context.view_layer.update()
+    bbox_min, bbox_max = scene_bbox()
+    offset = -(bbox_min + bbox_max) / 2
+    scene.matrix_world.translation += offset
+    
+    return scale, offset
+
+
+def main(arg):    
+    # Initialize context
+    if arg.object.endswith(".blend"):
+        delete_invisible_objects()
+    else:
+        init_scene()
+        load_object(arg.object)
+    print('[INFO] Scene initialized.')
+    
+    # Normalize scene
+    scale, offset = normalize_scene()
+    print('[INFO] Scene normalized.')
+    
+    # Start dumping
+    depsgraph = bpy.context.evaluated_depsgraph_get()
+    scene = bpy.context.scene
+    output = {
+        'objects': [],
+    }
+
+    # Dumping meshes
+    for obj in scene.objects:
+        if obj.type != 'MESH':
+            continue
+        
+        pack = {
+            "vertices": None,
+            "faces": None,
+        }
+        
+        eval_obj = obj.evaluated_get(depsgraph)
+        eval_mesh = eval_obj.to_mesh()
+        
+        bm = bmesh.new()
+        bm.from_mesh(eval_mesh)
+        bm.transform(obj.matrix_world)
+        bmesh.ops.triangulate(bm, faces=bm.faces)
+        bm.to_mesh(eval_mesh)
+        bm.free()
+                
+        pack["vertices"] = np.array([
+            v.co[:] for v in eval_mesh.vertices
+        ], dtype=np.float32)   # (N, 3)
+        
+        pack["faces"] = np.array([
+            [eval_mesh.loops[i].vertex_index for i in poly.loop_indices]
+            for poly in eval_mesh.polygons
+        ], dtype=np.int32)   # (F, 3)
+
+        output['objects'].append(pack)
+
+    # Save output
+    os.makedirs(os.path.dirname(arg.output_path), exist_ok=True)
+    with open(arg.output_path, 'wb') as f:
+        pickle.dump(output, f)
+    print('[INFO] Output saved to {}.'.format(arg.output_path))
+
+        
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Renders given obj file by rotation a camera around it.')
+    parser.add_argument('--object', type=str, help='Path to the 3D model file to be rendered.')
+    parser.add_argument('--output_path', type=str, default='/tmp', help='The path the output will be dumped to.')
+    argv = sys.argv[sys.argv.index("--") + 1:]
+    args = parser.parse_args(argv)
+
+    main(args)
+    
\ No newline at end of file
diff --git a/data_toolkit/blender_script/dump_pbr.py b/data_toolkit/blender_script/dump_pbr.py
new file mode 100644
index 0000000..302df2c
--- /dev/null
+++ b/data_toolkit/blender_script/dump_pbr.py
@@ -0,0 +1,485 @@
+import argparse, sys, os, math, io
+from typing import *
+import bpy
+import bmesh
+from mathutils import Vector, Matrix
+import numpy as np
+from PIL import Image
+import pickle
+
+
+"""=============== BLENDER ==============="""
+
+IMPORT_FUNCTIONS: Dict[str, Callable] = {
+    "obj": bpy.ops.import_scene.obj if bpy.app.version[0] < 4 else bpy.ops.wm.obj_import,
+    "glb": bpy.ops.import_scene.gltf,
+    "gltf": bpy.ops.import_scene.gltf,
+    "usd": bpy.ops.import_scene.usd,
+    "fbx": bpy.ops.import_scene.fbx,
+    "stl": bpy.ops.import_mesh.stl if bpy.app.version[0] < 4 else bpy.ops.wm.stl_import,
+    "usda": bpy.ops.import_scene.usda,
+    "dae": bpy.ops.wm.collada_import,
+    "ply": bpy.ops.import_mesh.ply if bpy.app.version[0] < 4 else bpy.ops.wm.ply_import,
+    "abc": bpy.ops.wm.alembic_import,
+    "blend": bpy.ops.wm.append,
+}
+
+
+def init_scene() -> None:
+    """Resets the scene to a clean state.
+
+    Returns:
+        None
+    """
+    # delete everything
+    for obj in bpy.data.objects:
+        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 = object_path.split(".")[-1].lower()
+    if file_extension is None:
+        raise ValueError(f"Unsupported file type: {object_path}")
+
+    if file_extension == "usdz":
+        # install usdz io package
+        dirname = os.path.dirname(os.path.realpath(__file__))
+        usdz_package = os.path.join(dirname, "io_scene_usdz.zip")
+        bpy.ops.preferences.addon_install(filepath=usdz_package)
+        # enable it
+        addon_name = "io_scene_usdz"
+        bpy.ops.preferences.addon_enable(module=addon_name)
+        # import the usdz
+        from io_scene_usdz.import_usdz import import_usdz
+
+        import_usdz(context, filepath=object_path, materials=True, animations=True)
+        return None
+
+    # load from existing import functions
+    import_function = IMPORT_FUNCTIONS[file_extension]
+
+    print(f"Loading object from {object_path}")
+    if file_extension == "blend":
+        import_function(directory=object_path, link=False)
+    elif file_extension in {"glb", "gltf"}:
+        import_function(filepath=object_path, merge_vertices=True, import_shading='NORMALS', bone_heuristic='TEMPERANCE')
+    else:
+        import_function(filepath=object_path)
+        
+        
+def delete_invisible_objects() -> None:
+    """Deletes all invisible objects in the scene.
+
+    Returns:
+        None
+    """
+    to_remove = []
+    for obj in bpy.context.scene.objects:
+        if obj.hide_viewport or obj.hide_render:
+            obj.hide_viewport = False
+            obj.hide_render = False
+            obj.hide_select = False
+            to_remove.append(obj)
+    for obj in to_remove:
+        bpy.data.objects.remove(obj, do_unlink=True)
+
+    # Delete invisible collections
+    invisible_collections = [col for col in bpy.data.collections if col.hide_viewport]
+    for col in invisible_collections:
+        bpy.data.collections.remove(col)
+      
+
+def scene_bbox() -> 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)
+
+    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
+    scene_meshes = [obj for obj in bpy.context.scene.objects.values() if isinstance(obj.data, bpy.types.Mesh)]
+    for obj in scene_meshes:
+        found = True
+        for coord in obj.bound_box:
+            coord = Vector(coord)
+            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 normalize_scene() -> Tuple[float, Vector]:
+    """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:
+        Tuple[float, Vector]: The scale factor and the offset applied to the scene.
+    """
+    scene_root_objects = [obj for obj in bpy.context.scene.objects.values() if not obj.parent]
+    if len(scene_root_objects) > 1:
+        # create an empty object to be used as a parent for all root objects
+        scene = bpy.data.objects.new("ParentEmpty", None)
+        bpy.context.scene.collection.objects.link(scene)
+
+        # parent all root objects to the empty object
+        for obj in scene_root_objects:
+            obj.parent = scene
+    else:
+        scene = scene_root_objects[0]
+
+    bbox_min, bbox_max = scene_bbox()
+    scale = 1 / max(bbox_max - bbox_min)
+    scene.scale = scene.scale * scale
+
+    # Apply scale to matrix_world.
+    bpy.context.view_layer.update()
+    bbox_min, bbox_max = scene_bbox()
+    offset = -(bbox_min + bbox_max) / 2
+    scene.matrix_world.translation += offset
+    
+    return scale, offset
+
+
+# =============== NODE TREE PARSING ===============
+
+def extract_image(tex_node, channels):
+        image = tex_node.image
+        pixels = np.array(image.pixels[:])
+        data = pixels.reshape(image.size[1], image.size[0], -1)
+        data = data[..., channels]
+
+        if data.dtype != np.uint8:
+            data = np.clip(data, 0.0, 1.0)
+            data = (data * 255).astype(np.uint8)
+
+        if len(data.shape) == 2:  # Single channel
+            pil_image = Image.fromarray(data, mode='L')
+        elif data.shape[2] == 3:
+            pil_image = Image.fromarray(data, mode='RGB')
+        elif data.shape[2] == 4:
+            pil_image = Image.fromarray(data, mode='RGBA')
+        else:
+            raise ValueError("Unsupported channel shape for image")
+
+        buffer = io.BytesIO()
+        pil_image.save(buffer, format='PNG')
+        png_bytes = buffer.getvalue()
+
+        return {
+            'image': png_bytes,
+            'interpolation': tex_node.interpolation,
+            'extension': tex_node.extension,
+        }
+
+
+def try_extract_image(link, expected_channel='RGB'):
+    """
+    Tries to extract an image from a texture node link.
+    Supported sub tree modes:
+      - RGB:
+        TEX_IMAGE ->
+      - R, G, B:
+        TEX_IMAGE -> SEPARATE_COLOR ->
+      - A:
+        TEX_IMAGE ->
+    """
+    assert expected_channel in ['RGB', 'R', 'G', 'B', 'A'], "Unsupported channel"
+
+    if expected_channel == 'RGB':
+        assert link.from_node.type == 'TEX_IMAGE', "Material is not supported"
+        assert link.from_socket.name == 'Color', "Material is not supported"
+        tex_node = link.from_node
+        return extract_image(tex_node, [0, 1, 2])
+
+    if expected_channel in ['R', 'G', 'B']:
+        socket_name = {
+            'R': 'Red',
+            'G': 'Green',
+            'B': 'Blue',
+        }[expected_channel]
+        assert link.from_node.type == 'SEPARATE_COLOR' and link.from_node.mode == 'RGB', \
+            f"Material is not supported, {link.from_node.type}, {link.from_node.mode}"
+        assert link.from_socket.name == socket_name, "Material is not supported"
+        sep_node = link.from_node
+        assert sep_node.inputs[0].is_linked and sep_node.inputs[0].links[0].from_node.type == 'TEX_IMAGE', \
+            "Material is not supported"
+        assert sep_node.inputs[0].links[0].from_socket.name == 'Color', "Material is not supported"
+        tex_node = sep_node.inputs[0].links[0].from_node
+        channel_index = {
+            'R': 0,
+            'G': 1,
+            'B': 2,
+        }[expected_channel]
+        return extract_image(tex_node, channel_index)
+
+    if expected_channel == 'A':
+        assert link.from_node.type == 'TEX_IMAGE', "Material is not supported"
+        assert link.from_socket.name == 'Alpha', "Material is not supported"
+        tex_node = link.from_node
+        return extract_image(tex_node, 3)
+
+
+def try_extract_factor(link, mode='color'):
+    """
+    Tries to extract a factor from a math node link.
+    Supported sub tree modes:
+      - color:
+       ANY -> MIX(MULTIPLY) ->
+      - scalar:
+       ANY -> MATH(MULTIPLY) ->
+    """
+    assert mode in ['color','scalar'], "Unsupported mode"
+
+    if mode == 'color':
+        if link.from_node.type == 'MIX':
+            mix_node = link.from_node
+            assert mix_node.data_type == 'RGBA' and mix_node.blend_type == 'MULTIPLY', f"Material is not supported, {mix_node.data_type}, {mix_node.blend_type}"
+            assert not mix_node.inputs['Factor'].is_linked and mix_node.inputs['Factor'].default_value == 1.0, \
+                "Material is not supported"
+            if mix_node.inputs['A'].is_linked:
+                assert not mix_node.inputs['B'].is_linked, "Material is not supported"
+                return (list(mix_node.inputs['B'].default_value)[:3], mix_node.inputs['A'].links[0])
+            else:
+                assert not mix_node.inputs['A'].is_linked, "Material is not supported"
+                assert mix_node.inputs['B'].is_linked, "Material is not supported"
+                return (list(mix_node.inputs['A'].default_value)[:3], mix_node.inputs['B'].links[0])
+        return ([1.0, 1.0, 1.0], link)
+
+    if mode =='scalar':
+        if link.from_node.type == 'MATH':
+            math_node = link.from_node
+            assert math_node.operation == 'MULTIPLY', "Material is not supported"
+            assert math_node.inputs[0].is_linked, "Material is not supported"
+            assert not math_node.inputs[1].is_linked, "Material is not supported"
+            return (math_node.inputs[1].default_value, math_node.inputs[0].links[0])
+        return (1.0, link)
+
+
+def try_extract_image_with_factor(link, expected_channel='RGB'):
+    """
+    Tries to extract an image and a factor from a texture node link.
+    """
+    factor, link = try_extract_factor(link, 'color' if expected_channel in ['RGB'] else 'scalar')
+    image = try_extract_image(link, expected_channel)
+    return (factor, image)
+
+
+def main(arg):    
+    # Initialize context
+    if arg.object.endswith(".blend"):
+        delete_invisible_objects()
+    else:
+        init_scene()
+        load_object(arg.object)
+    print('[INFO] Scene initialized.')
+    
+    # Normalize scene
+    scale, offset = normalize_scene()
+    print('[INFO] Scene normalized.')
+    
+    # Start dumping
+    depsgraph = bpy.context.evaluated_depsgraph_get()
+    scene = bpy.context.scene
+    output = {
+        'materials': [],
+        'objects': [],
+    }
+
+    # Dumping materials
+    for mat in bpy.data.materials:
+        assert mat.use_nodes == True, "Material is not supported"
+
+        pack = {
+            "baseColorFactor": [1.0, 1.0, 1.0],
+            "alphaFactor": 1.0,
+            "metallicFactor": 1.0,
+            "roughnessFactor": 1.0,
+            "alphaMode": "OPAQUE",
+            "alphaCutoff": 0.5,
+            "baseColorTexture": None,
+            "alphaTexture": None,
+            "metallicTexture": None,
+            "roughnessTexture": None,
+        }
+
+        try:
+            principled_node = mat.node_tree.nodes.get('Principled BSDF')
+            assert principled_node is not None, "Material is not supported"
+
+            # Handle base color
+            if not principled_node.inputs['Base Color'].is_linked:
+                pack["baseColorFactor"] = list(principled_node.inputs['Base Color'].default_value)
+            else:
+                link = principled_node.inputs['Base Color'].links[0]
+                if link.from_node.type == 'RGB':
+                    pack["baseColorFactor"] = list(link.from_node.outputs[0].default_value)
+                else:
+                    factor, image = try_extract_image_with_factor(link, 'RGB')
+                    pack["baseColorFactor"] = factor
+                    pack["baseColorTexture"] = image
+
+            # Handle alpha
+            if not principled_node.inputs['Alpha'].is_linked:
+                pack["alphaFactor"] = principled_node.inputs['Alpha'].default_value
+                if pack["alphaFactor"] < 1.0:
+                    pack["alphaMode"] = "BLEND"
+            else:
+                link = principled_node.inputs['Alpha'].links[0]
+                node = link.from_node
+                if node.type == 'VALUE':
+                    pack["alphaFactor"] = node.outputs[0].default_value
+                    if pack["alphaFactor"] < 1.0:
+                        pack["alphaMode"] = "BLEND"
+                else:
+                    pack["alphaMode"] = "BLEND"
+                    if node.type == 'MATH':
+                        if node.operation == 'ROUND':
+                            assert node.inputs[0].is_linked, "Material is not supported"
+                            pack["alphaMode"] = "MASK"
+                            link = node.inputs[0].links[0]
+                        elif node.operation == 'SUBTRACT':
+                            assert node.inputs[0].default_value == 1.0 and \
+                                node.inputs[1].is_linked and \
+                                node.inputs[1].links[0].from_node.type == 'MATH' and \
+                                node.inputs[1].links[0].from_node.operation == 'LESS_THAN', \
+                                "Material is not supported"
+                            assert node.inputs[1].links[0].from_node.inputs[0].is_linked, "Material is not supported"
+                            pack["alphaMode"] = "MASK"
+                            pack["alphaCutoff"] = node.inputs[1].links[0].from_node.inputs[1].default_value
+                            link = node.inputs[1].links[0].from_node.inputs[0].links[0]
+                    factor, image = try_extract_image_with_factor(link, 'A')
+                    pack["alphaFactor"] = factor
+                    pack["alphaTexture"] = image
+
+            # Handle metallic
+            if not principled_node.inputs['Metallic'].is_linked:
+                pack["metallicFactor"] = principled_node.inputs['Metallic'].default_value
+            else:
+                link = principled_node.inputs['Metallic'].links[0]
+                node = link.from_node
+                if node.type == 'VALUE':
+                    pack["metallicFactor"] = node.outputs[0].default_value
+                else:
+                    factor, image = try_extract_image_with_factor(link, 'B')
+                    pack["metallicFactor"] = factor
+                    pack["metallicTexture"] = image
+
+            # Handle roughness
+            if not principled_node.inputs['Roughness'].is_linked:
+                pack["roughnessFactor"] = principled_node.inputs['Roughness'].default_value
+            else:
+                link = principled_node.inputs['Roughness'].links[0]
+                node = link.from_node
+                if node.type == 'VALUE':
+                    pack["roughnessFactor"] = node.outputs[0].default_value
+                else:
+                    factor, image = try_extract_image_with_factor(link, 'G')
+                    pack["roughnessFactor"] = factor
+                    pack["roughnessTexture"] = image
+
+            output['materials'].append(pack)
+        except:
+            with open(arg.output_path + '_error.txt', 'w') as f:
+                f.write(str([[n.name] for n in mat.node_tree.nodes]))
+            raise RuntimeError("Material is not supported")
+
+    # Dumping meshes
+    for obj in scene.objects:
+        if obj.type != 'MESH':
+            continue
+        
+        pack = {
+            "vertices": None,
+            "faces": None,
+            "uvs": None,
+            "matIDs": None,
+        }
+        
+        eval_obj = obj.evaluated_get(depsgraph)
+        eval_mesh = eval_obj.to_mesh()
+        
+        bm = bmesh.new()
+        bm.from_mesh(eval_mesh)
+        bm.transform(obj.matrix_world)
+        bmesh.ops.triangulate(bm, faces=bm.faces)
+        bm.to_mesh(eval_mesh)
+        bm.free()
+                
+        pack["vertices"] = np.array([
+            v.co[:] for v in eval_mesh.vertices
+        ], dtype=np.float32)   # (N, 3)
+        
+        pack["faces"] = np.array([
+            [eval_mesh.loops[i].vertex_index for i in poly.loop_indices]
+            for poly in eval_mesh.polygons
+        ], dtype=np.int32)   # (F, 3)
+        
+        pack["normals"] = np.array([
+            [eval_mesh.loops[i].normal for i in poly.loop_indices]
+            for poly in eval_mesh.polygons
+        ], dtype=np.float32)  # (F, 3, 3)
+        
+        if eval_mesh.uv_layers.active is not None:
+            pack["uvs"] = np.array([
+                [eval_mesh.uv_layers.active.data[i].uv for i in poly.loop_indices]
+                for poly in eval_mesh.polygons
+            ], dtype=np.float32)  # (F, 3, 2)
+
+        pack["mat_ids"] = np.array([
+            bpy.data.materials.find(obj.material_slots[poly.material_index].name)
+            if len(obj.material_slots) > 0 and obj.material_slots[poly.material_index].material is not None else -1
+            for poly in eval_mesh.polygons
+        ], dtype=np.int32)
+
+        output['objects'].append(pack)
+
+    # Save output
+    os.makedirs(os.path.dirname(arg.output_path), exist_ok=True)
+    with open(arg.output_path, 'wb') as f:
+        pickle.dump(output, f)
+    print('[INFO] Output saved to {}.'.format(arg.output_path))
+
+        
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Renders given obj file by rotation a camera around it.')
+    parser.add_argument('--object', type=str, help='Path to the 3D model file to be rendered.')
+    parser.add_argument('--output_path', type=str, default='/tmp', help='The path the output will be dumped to.')
+    argv = sys.argv[sys.argv.index("--") + 1:]
+    args = parser.parse_args(argv)
+
+    main(args)
+    
\ No newline at end of file
diff --git a/data_toolkit/blender_script/install_pillow.py b/data_toolkit/blender_script/install_pillow.py
new file mode 100644
index 0000000..b0b0b0c
--- /dev/null
+++ b/data_toolkit/blender_script/install_pillow.py
@@ -0,0 +1,6 @@
+import subprocess
+import sys
+import ensurepip
+
+ensurepip.bootstrap()
+subprocess.check_call([sys.executable, "-m", "pip", "install", "Pillow"])
\ No newline at end of file
diff --git a/data_toolkit/blender_script/io_scene_usdz.zip b/data_toolkit/blender_script/io_scene_usdz.zip
new file mode 100644
index 0000000..f2e2fdb
Binary files /dev/null and b/data_toolkit/blender_script/io_scene_usdz.zip differ
diff --git a/data_toolkit/blender_script/render_cond.py b/data_toolkit/blender_script/render_cond.py
new file mode 100644
index 0000000..b645418
--- /dev/null
+++ b/data_toolkit/blender_script/render_cond.py
@@ -0,0 +1,529 @@
+import argparse, sys, os, math, re, glob
+from typing import *
+import bpy
+from mathutils import Vector, Matrix
+import numpy as np
+import json
+import glob
+from PIL import Image
+
+
+"""=============== BLENDER ==============="""
+
+IMPORT_FUNCTIONS: Dict[str, Callable] = {
+    "obj": bpy.ops.import_scene.obj,
+    "glb": bpy.ops.import_scene.gltf,
+    "gltf": bpy.ops.import_scene.gltf,
+    "usd": bpy.ops.import_scene.usd,
+    "fbx": bpy.ops.import_scene.fbx,
+    "stl": bpy.ops.import_mesh.stl,
+    "usda": bpy.ops.import_scene.usda,
+    "dae": bpy.ops.wm.collada_import,
+    "ply": bpy.ops.import_mesh.ply,
+    "abc": bpy.ops.wm.alembic_import,
+    "blend": bpy.ops.wm.append,
+}
+
+EXT = {
+    'PNG': 'png',
+    'JPEG': 'jpg',
+    'OPEN_EXR': 'exr',
+    'TIFF': 'tiff',
+    'BMP': 'bmp',
+    'HDR': 'hdr',
+    'TARGA': 'tga'
+}
+
+
+def init_render(engine='CYCLES', resolution=512):
+    bpy.context.scene.render.engine = engine
+    bpy.context.scene.render.resolution_x = resolution
+    bpy.context.scene.render.resolution_y = resolution
+    bpy.context.scene.render.resolution_percentage = 100
+    bpy.context.scene.render.image_settings.file_format = 'PNG'
+    bpy.context.scene.render.image_settings.color_mode = 'RGBA'
+    bpy.context.scene.render.film_transparent = True
+    
+    bpy.context.scene.cycles.device = 'GPU'
+    bpy.context.scene.cycles.samples = 32
+    bpy.context.scene.cycles.filter_type = 'BOX'
+    bpy.context.scene.cycles.filter_width = 1
+    bpy.context.scene.cycles.diffuse_bounces = 1
+    bpy.context.scene.cycles.glossy_bounces = 1
+    bpy.context.scene.cycles.transparent_max_bounces = 3
+    bpy.context.scene.cycles.transmission_bounces = 3
+    bpy.context.scene.cycles.use_denoising = True
+        
+    bpy.context.preferences.addons['cycles'].preferences.get_devices()
+    bpy.context.preferences.addons['cycles'].preferences.compute_device_type = 'CUDA'
+    
+
+def init_scene() -> None:
+    """Resets the scene to a clean state.
+
+    Returns:
+        None
+    """
+    # delete everything
+    for obj in bpy.data.objects:
+        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 init_camera():
+    cam = bpy.data.objects.new('Camera', bpy.data.cameras.new('Camera'))
+    bpy.context.collection.objects.link(cam)
+    bpy.context.scene.camera = cam
+    cam.data.sensor_height = cam.data.sensor_width = 32
+    cam_constraint = cam.constraints.new(type='TRACK_TO')
+    cam_constraint.track_axis = 'TRACK_NEGATIVE_Z'
+    cam_constraint.up_axis = 'UP_Y'
+    cam_empty = bpy.data.objects.new("Empty", None)
+    cam_empty.location = (0, 0, 0)
+    bpy.context.scene.collection.objects.link(cam_empty)
+    cam_constraint.target = cam_empty
+    return cam
+
+
+def init_uniform_lighting():
+    # Clear existing lights
+    bpy.ops.object.select_all(action="DESELECT")
+    bpy.ops.object.select_by_type(type="LIGHT")
+    bpy.ops.object.delete()
+    
+    # Create environment light
+    if bpy.context.scene.world is None:
+        world = bpy.data.worlds.new("World")
+        bpy.context.scene.world = world
+    else:
+        world = bpy.context.scene.world
+        
+    # Enabling nodes
+    world.use_nodes = True
+    node_tree = world.node_tree
+    nodes = node_tree.nodes
+    links = node_tree.links
+    
+    # Remove default nodes
+    for node in nodes:
+        nodes.remove(node)
+
+    # Create background node
+    bg_node = nodes.new(type="ShaderNodeBackground")
+    bg_node.inputs["Color"].default_value = (1.0, 1.0, 1.0, 1.0)
+    bg_node.inputs["Strength"].default_value = 1.0
+    output_node = nodes.new(type="ShaderNodeOutputWorld")
+    links.new(bg_node.outputs["Background"], output_node.inputs["Surface"])
+
+
+def init_random_lighting(camera_dir: np.ndarray) -> None:
+    # Clear existing lights
+    bpy.ops.object.select_all(action="DESELECT")
+    bpy.ops.object.select_by_type(type="LIGHT")
+    bpy.ops.object.delete()
+    
+    # Create environment light
+    if bpy.context.scene.world is None:
+        world = bpy.data.worlds.new("World")
+        bpy.context.scene.world = world
+    else:
+        world = bpy.context.scene.world
+        
+    # Enabling nodes
+    world.use_nodes = True
+    node_tree = world.node_tree
+    nodes = node_tree.nodes
+    links = node_tree.links
+    
+    # Remove default nodes
+    for node in nodes:
+        nodes.remove(node)
+    
+    # Random place lights
+    num_lights = np.random.randint(1, 4)
+    total_strength = 1.5
+    for i in range(num_lights):
+        new_light = bpy.data.objects.new(f"Light_{i}", bpy.data.lights.new(f"Light_{i}", type="POINT"))
+        bpy.context.collection.objects.link(new_light)
+        
+        new_light_distance = 1 / np.random.uniform(1/100, 1/10)
+        new_light_dir = np.random.randn(3)
+        new_light_dir[2] += 0.6
+        new_light_dir = new_light_dir / np.linalg.norm(new_light_dir)
+        new_light_location = new_light_dir * new_light_distance
+        new_light_camera_strength_ratio = max(np.sum(camera_dir * new_light_dir) * 0.5 + 0.5, 0)
+        new_light_max_energy = total_strength / (np.sum(camera_dir * new_light_dir) * 0.45 + 0.55)
+        new_light_strength = np.sqrt(np.random.uniform(0.01, 1)) * new_light_max_energy
+        new_light_camera_strength = new_light_camera_strength_ratio * new_light_strength
+        total_strength -= new_light_camera_strength
+        
+        new_light.location = (new_light_location[0], new_light_location[1], new_light_location[2])
+        new_light.data.color = (1.0, 1.0, 1.0)
+        new_light.data.energy = new_light_strength * new_light_distance**2 * 31.4
+        new_light.data.shadow_soft_size = np.random.uniform(0.1, 0.1 * new_light_distance)
+        
+    # Create background node
+    bg_node = nodes.new(type="ShaderNodeBackground")
+    bg_node.inputs["Color"].default_value = (1.0, 1.0, 1.0, 1.0)
+    bg_node.inputs["Strength"].default_value = total_strength
+    output_node = nodes.new(type="ShaderNodeOutputWorld")
+    links.new(bg_node.outputs["Background"], output_node.inputs["Surface"])
+
+
+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 = object_path.split(".")[-1].lower()
+    if file_extension is None:
+        raise ValueError(f"Unsupported file type: {object_path}")
+
+    if file_extension == "usdz":
+        # install usdz io package
+        dirname = os.path.dirname(os.path.realpath(__file__))
+        usdz_package = os.path.join(dirname, "io_scene_usdz.zip")
+        bpy.ops.preferences.addon_install(filepath=usdz_package)
+        # enable it
+        addon_name = "io_scene_usdz"
+        bpy.ops.preferences.addon_enable(module=addon_name)
+        # import the usdz
+        from io_scene_usdz.import_usdz import import_usdz
+
+        import_usdz(context, filepath=object_path, materials=True, animations=True)
+        return None
+
+    # load from existing import functions
+    import_function = IMPORT_FUNCTIONS[file_extension]
+
+    print(f"Loading object from {object_path}")
+    if file_extension == "blend":
+        import_function(directory=object_path, link=False)
+    elif file_extension in {"glb", "gltf"}:
+        import_function(filepath=object_path, merge_vertices=True, import_shading='NORMALS')
+    else:
+        import_function(filepath=object_path)
+        
+
+def delete_invisible_objects() -> None:
+    """Deletes all invisible objects in the scene.
+
+    Returns:
+        None
+    """
+    # bpy.ops.object.mode_set(mode="OBJECT")
+    bpy.ops.object.select_all(action="DESELECT")
+    for obj in bpy.context.scene.objects:
+        if obj.hide_viewport or obj.hide_render:
+            obj.hide_viewport = False
+            obj.hide_render = False
+            obj.hide_select = False
+            obj.select_set(True)
+    bpy.ops.object.delete()
+
+    # Delete invisible collections
+    invisible_collections = [col for col in bpy.data.collections if col.hide_viewport]
+    for col in invisible_collections:
+        bpy.data.collections.remove(col)
+
+
+def unhide_all_objects() -> None:
+    """Unhides all objects in the scene.
+
+    Returns:
+        None
+    """
+    for obj in bpy.context.scene.objects:
+        obj.hide_set(False)
+        
+        
+def convert_to_meshes() -> None:
+    """Converts all objects in the scene to meshes.
+
+    Returns:
+        None
+    """
+    bpy.ops.object.select_all(action="DESELECT")
+    bpy.context.view_layer.objects.active = [obj for obj in bpy.context.scene.objects if obj.type == "MESH"][0]
+    for obj in bpy.context.scene.objects:
+        obj.select_set(True)
+    bpy.ops.object.convert(target="MESH")
+    
+        
+def triangulate_meshes() -> None:
+    """Triangulates all meshes in the scene.
+
+    Returns:
+        None
+    """
+    bpy.ops.object.select_all(action="DESELECT")
+    objs = [obj for obj in bpy.context.scene.objects if obj.type == "MESH"]
+    bpy.context.view_layer.objects.active = objs[0]
+    for obj in objs:
+        obj.select_set(True)
+    bpy.ops.object.mode_set(mode="EDIT")
+    bpy.ops.mesh.reveal()
+    bpy.ops.mesh.select_all(action="SELECT")
+    bpy.ops.mesh.quads_convert_to_tris(quad_method="BEAUTY", ngon_method="BEAUTY")
+    bpy.ops.object.mode_set(mode="OBJECT")
+    bpy.ops.object.select_all(action="DESELECT")
+    
+
+def scene_bbox() -> 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)
+
+    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
+    scene_meshes = [obj for obj in bpy.context.scene.objects.values() if isinstance(obj.data, bpy.types.Mesh)]
+    for obj in scene_meshes:
+        found = True
+        for coord in obj.bound_box:
+            coord = Vector(coord)
+            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 normalize_scene() -> Tuple[float, Vector]:
+    """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:
+        Tuple[float, Vector]: The scale factor and the offset applied to the scene.
+    """
+    scene_root_objects = [obj for obj in bpy.context.scene.objects.values() if not obj.parent]
+    if len(scene_root_objects) > 1:
+        # create an empty object to be used as a parent for all root objects
+        scene = bpy.data.objects.new("ParentEmpty", None)
+        bpy.context.scene.collection.objects.link(scene)
+
+        # parent all root objects to the empty object
+        for obj in scene_root_objects:
+            obj.parent = scene
+    else:
+        scene = scene_root_objects[0]
+
+    bbox_min, bbox_max = scene_bbox()
+    scale = 1 / max(bbox_max - bbox_min)
+    scene.scale = scene.scale * scale
+
+    # Apply scale to matrix_world.
+    bpy.context.view_layer.update()
+    bbox_min, bbox_max = scene_bbox()
+    offset = -(bbox_min + bbox_max) / 2
+    scene.matrix_world.translation += offset
+    bpy.ops.object.select_all(action="DESELECT")
+    
+    return scale, offset
+
+
+def get_transform_matrix(obj: bpy.types.Object) -> list:
+    pos, rt, _ = obj.matrix_world.decompose()
+    rt = rt.to_matrix()
+    matrix = []
+    for ii in range(3):
+        a = []
+        for jj in range(3):
+            a.append(rt[ii][jj])
+        a.append(pos[ii])
+        matrix.append(a)
+    matrix.append([0, 0, 0, 1])
+    return matrix
+
+
+def check_mask_boundary_distance(image_path: str, threshold: int = 0) -> Tuple[bool, bool, int]:
+    """Check the rendered object's mask distance to image boundary.
+    
+    Args:
+        image_path: Path to the rendered PNG image with alpha channel.
+        threshold: Alpha value threshold to consider as valid mask (0-255).
+        
+    Returns:
+        Tuple of (touches_boundary, too_far_from_boundary, min_distance):
+        - touches_boundary: True if the mask touches any boundary
+        - too_far_from_boundary: True if the mask is too far from all boundaries (>80 pixels)
+        - min_distance: Minimum distance from mask to any boundary
+    """
+    img = Image.open(image_path)
+    if img.mode != 'RGBA':
+        return False, False, 0
+    
+    # Get alpha channel
+    alpha = np.array(img)[:, :, 3]
+    h, w = alpha.shape
+    
+    # Find all pixels with alpha > threshold (mask pixels)
+    mask_pixels = np.where(alpha > threshold)
+    
+    if len(mask_pixels[0]) == 0:
+        # No valid mask pixels
+        return False, True, max(h, w)
+    
+    # Calculate distances to each boundary
+    min_row = np.min(mask_pixels[0])  # Distance to top edge
+    max_row = np.max(mask_pixels[0])  # Distance to bottom edge
+    min_col = np.min(mask_pixels[1])  # Distance to left edge
+    max_col = np.max(mask_pixels[1])  # Distance to right edge
+    
+    dist_top = min_row
+    dist_bottom = (h - 1) - max_row
+    dist_left = min_col
+    dist_right = (w - 1) - max_col
+    
+    min_distance = min(dist_top, dist_bottom, dist_left, dist_right)
+    
+    # Check if touches boundary (distance <= 0)
+    touches_boundary = min_distance <= 0
+    
+    # Check if too far from boundary (distance > 130 pixels)
+    too_far = min_distance > 130
+    
+    return touches_boundary, too_far, min_distance
+
+
+def main(arg):
+    if arg.object.endswith(".blend"):
+        delete_invisible_objects()
+    else:
+        init_scene()
+        load_object(arg.object)
+    print('[INFO] Scene initialized.')
+    
+    # normalize scene
+    scale, offset = normalize_scene()
+    print('[INFO] Scene normalized.')
+    
+    # Initialize camera and lighting
+    cam = init_camera()
+    init_uniform_lighting()
+    print('[INFO] Camera and lighting initialized.')
+        
+    # ============= Render conditional views =============
+    init_render(engine=arg.engine, resolution=arg.cond_resolution)
+    # Create a list of views
+    to_export = {
+        "aabb": [[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
+        "scale": scale,
+        "offset": [offset.x, offset.y, offset.z],
+        "frames": []
+    }
+    views = json.loads(arg.cond_views)
+    
+    # Parameters for boundary check and radius adjustment
+    max_retry = 10  # Maximum number of retries per view
+    radius_increase_factor = 1.1  # Increase radius by 10% when too close to boundary
+    radius_decrease_factor = 0.9  # Decrease radius by 10% when too far from boundary
+    min_boundary_distance = 130  # Minimum distance to boundary in pixels
+    
+    for i, view in enumerate(views):
+        current_radius = view['radius']
+        retry_count = 0
+        
+        while retry_count < max_retry:
+            cam_dir = np.array([
+                np.cos(view['yaw']) * np.cos(view['pitch']),
+                np.sin(view['yaw']) * np.cos(view['pitch']),
+                np.sin(view['pitch'])
+            ])
+            init_random_lighting(cam_dir)
+            cam.location = (
+                current_radius * cam_dir[0],
+                current_radius * cam_dir[1],
+                current_radius * cam_dir[2]
+            )
+            cam.data.lens = 16 / np.tan(view['fov'] / 2)
+            
+            output_path = os.path.join(arg.cond_output_folder, f'{i:03d}.png')
+            bpy.context.scene.render.filepath = output_path
+                
+            # Render the scene
+            bpy.ops.render.render(write_still=True)
+            bpy.context.view_layer.update()
+            
+            # Check mask boundary distance
+            touches_boundary, too_far, min_dist = check_mask_boundary_distance(output_path)
+            
+            if touches_boundary:
+                # Object is too close to boundary, increase radius
+                retry_count += 1
+                old_radius = current_radius
+                current_radius *= radius_increase_factor
+                print(f'[WARNING] View {i}: Mask touches boundary (dist={min_dist}px). Increasing radius from {old_radius:.4f} to {current_radius:.4f} (retry {retry_count}/{max_retry})')
+            elif too_far:
+                # Object is too far from boundary (>80px), decrease radius
+                retry_count += 1
+                old_radius = current_radius
+                current_radius *= radius_decrease_factor
+                print(f'[INFO] View {i}: Mask too far from boundary (dist={min_dist}px > {min_boundary_distance}px). Decreasing radius from {old_radius:.4f} to {current_radius:.4f} (retry {retry_count}/{max_retry})')
+            else:
+                # Good distance, stop retrying
+                if retry_count > 0:
+                    print(f'[INFO] View {i}: Mask boundary distance OK (dist={min_dist}px) after {retry_count} retries. Final radius: {current_radius:.4f}')
+                else:
+                    print(f'[INFO] View {i}: Mask boundary distance OK (dist={min_dist}px). Radius: {current_radius:.4f}')
+                break
+        
+        if retry_count >= max_retry:
+            print(f'[WARNING] View {i}: Max retries reached. Using final radius: {current_radius:.4f} (dist={min_dist}px)')
+            
+        # Save camera parameters (with potentially updated radius)
+        metadata = {
+            "file_path": f'{i:03d}.png',
+            "camera_angle_x": view['fov'],
+            "transform_matrix": get_transform_matrix(cam),
+            "radius": current_radius,  # Save the actual radius used
+            "original_radius": view['radius'],  # Save original for reference
+            "retries": retry_count
+        }
+        to_export["frames"].append(metadata)
+    
+    # Save the camera parameters
+    with open(os.path.join(arg.cond_output_folder, 'transforms.json'), 'w') as f:
+        json.dump(to_export, f, indent=4)
+
+        
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Renders given obj file by rotation a camera around it.')
+    parser.add_argument('--object', type=str, help='Path to the 3D model file to be rendered.')
+    parser.add_argument('--cond_views', type=str, help='JSON string of views. Contains a list of {yaw, pitch, radius, fov} object.')
+    parser.add_argument('--cond_output_folder', type=str, default='/tmp', help='The path the output will be dumped to.')
+    parser.add_argument('--cond_resolution', type=int, default=1024, help='Resolution of the conditional images.')
+    parser.add_argument('--engine', type=str, default='CYCLES', help='Blender internal engine for rendering. E.g. CYCLES, BLENDER_EEVEE, ...')
+    argv = sys.argv[sys.argv.index("--") + 1:]
+    args = parser.parse_args(argv)
+
+    main(args)
+    
\ No newline at end of file
diff --git a/data_toolkit/build_metadata.py b/data_toolkit/build_metadata.py
new file mode 100644
index 0000000..c0ee863
--- /dev/null
+++ b/data_toolkit/build_metadata.py
@@ -0,0 +1,410 @@
+import os
+import shutil
+import sys
+import time
+import glob
+import importlib
+import argparse
+import pandas as pd
+from easydict import EasyDict as edict
+
+
+def update_metadata(path, opt):
+    if not os.path.exists(path):
+        return None
+    timestamp = str(int(time.time()))
+    os.makedirs(os.path.join(path, 'merged_records'), exist_ok=True)
+    os.makedirs(os.path.join(path, 'new_records'), exist_ok=True)
+    if opt.from_merged_records:
+        df_files = [f for f in os.listdir(os.path.join(path, 'merged_records')) if f.endswith('.csv')]
+        df_files = [f for f in df_files if int(f.split('_')[0]) >= opt.record_start]
+    else:
+        df_files = [f for f in os.listdir(os.path.join(path, 'new_records')) if f.startswith('part_') and f.endswith('.csv')]
+    df_parts = []
+    for f in df_files:
+        try:
+            df_parts.append(pd.read_csv(os.path.join(path, 'new_records', f)))
+        except Exception as e:
+            print(f"Failed to read {f}: {e}")
+    if len(df_parts) > 0:
+        if os.path.exists(os.path.join(path, 'metadata.csv')):
+            metadata = pd.read_csv(os.path.join(path, 'metadata.csv'))
+        else:
+            columns = df_parts[0].columns
+            metadata = pd.DataFrame(columns=columns)
+        metadata.set_index('sha256', inplace=True)
+        if metadata.index.duplicated().any():
+            metadata = metadata.groupby(level=0).first()
+        for df_part in df_parts:
+            if 'sha256' in df_part.columns:
+                df_part.set_index('sha256', inplace=True)
+                if df_part.index.duplicated().any():
+                    df_part = df_part.groupby(level=0).first()
+                metadata = df_part.combine_first(metadata)
+        metadata.to_csv(os.path.join(path, 'metadata.csv'))
+        for f in df_files:
+            shutil.move(os.path.join(path, 'new_records', f), os.path.join(path, 'merged_records', f'{timestamp}_{f}'))
+        return metadata
+    else:
+        if os.path.exists(os.path.join(path, 'metadata.csv')):
+            return pd.read_csv(os.path.join(path, 'metadata.csv'))
+    return None
+
+
+def build_downloaded_metadata_from_files(raw_root, global_metadata):
+    """Scan local files under raw_root to build download metadata.
+    
+    Walks through raw_root to find downloaded 3D files (.glb, .obj, .fbx, .usdz, .gltf, .zip),
+    matches them against global_metadata via file_identifier (uid extracted from URL) to recover
+    the sha256 -> local_path mapping.
+    """
+    extensions = ('.glb', '.obj', '.fbx', '.usdz', '.gltf', '.zip')
+    
+    # Build uid -> sha256 mapping from global metadata
+    uid_to_sha256 = {}
+    if 'file_identifier' in global_metadata.columns:
+        for _, row in global_metadata.iterrows():
+            uid = str(row['file_identifier']).split('/')[-1]
+            uid_to_sha256[uid] = row['sha256']
+    
+    # Scan files
+    records = []
+    for dirpath, dirnames, filenames in os.walk(raw_root):
+        for fname in filenames:
+            if not fname.lower().endswith(extensions):
+                continue
+            uid = os.path.splitext(fname)[0]
+            sha256 = uid_to_sha256.get(uid)
+            if sha256 is not None:
+                full_path = os.path.join(dirpath, fname)
+                # Store path relative to parent of raw_root (i.e. download_root)
+                rel_path = os.path.relpath(full_path, os.path.dirname(raw_root))
+                records.append({'sha256': sha256, 'local_path': rel_path})
+    
+    if len(records) == 0:
+        return None
+    
+    df = pd.DataFrame(records).set_index('sha256')
+    print(f'  [from_file] Found {len(df)} downloaded files under {raw_root}')
+    
+    # Save as metadata.csv under raw_root
+    os.makedirs(raw_root, exist_ok=True)
+    df.to_csv(os.path.join(raw_root, 'metadata.csv'))
+    return df
+
+
+# Check if directory is a multi-view directory (ending with _view or _view_fix)
+def _is_view_dir(dirname):
+    return dirname.endswith('_view') or dirname.endswith('_view_fix')
+
+
+if __name__ == '__main__':
+    dataset_utils = importlib.import_module(f'datasets.{sys.argv[1]}')
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--root', type=str, required=True,
+                        help='Directory to save the metadata')
+    parser.add_argument('--download_root', type=str, default=None,
+                        help='Directory to save the downloaded files')
+    parser.add_argument('--thumbnail_root', type=str, default=None,
+                        help='Directory to save the thumbnail files')
+    parser.add_argument('--render_cond_root', type=str, default=None,
+                        help='Directory to save the render condition files')
+    parser.add_argument('--mesh_dump_root', type=str, default=None,
+                        help='Directory to save the mesh files')
+    parser.add_argument('--pbr_dump_root', type=str, default=None,
+                        help='Directory to save the pbr files')
+    parser.add_argument('--dual_grid_root', type=str, default=None,
+                        help='Directory to save the dual grid files')
+    parser.add_argument('--pbr_voxel_root', type=str, default=None,
+                        help='Directory to save the pbr voxel files')
+    parser.add_argument('--ss_latent_root', type=str, default=None,
+                        help='Directory to save the sparse structure latent files')
+    parser.add_argument('--shape_latent_root', type=str, default=None,
+                        help='Directory to save the shape latent files')
+    parser.add_argument('--pbr_latent_root', type=str, default=None,
+                        help='Directory to save the pbr latent files')
+    parser.add_argument('--field', type=str, default='all',
+                        help='Fields to process, separated by commas')
+    parser.add_argument('--from_file', action='store_true',
+                        help='Build metadata from file instead of from records of processings.' +
+                             'Useful when some processing fail to generate records but file already exists.')
+    parser.add_argument('--from_merged_records', action='store_true',
+                        help='Build metadata from merged records')
+    parser.add_argument('--record_start', type=int)
+    parser.add_argument('--rebuild', action='store_true',
+                        help='Rebuild metadata from scratch, ignore existing metadata.')
+    dataset_utils.add_args(parser)
+    opt = parser.parse_args(sys.argv[2:])
+    opt = edict(vars(opt))
+    opt.download_root = opt.download_root or opt.root
+    opt.thumbnail_root = opt.thumbnail_root or opt.root
+    opt.render_cond_root = opt.render_cond_root or opt.root
+    opt.mesh_dump_root = opt.mesh_dump_root or opt.root
+    opt.pbr_dump_root = opt.pbr_dump_root or opt.root
+    opt.dual_grid_root = opt.dual_grid_root or opt.root
+    opt.pbr_voxel_root = opt.pbr_voxel_root or opt.root
+    opt.ss_latent_root = opt.ss_latent_root or opt.root
+    opt.shape_latent_root = opt.shape_latent_root or opt.root
+    opt.pbr_latent_root = opt.pbr_latent_root or opt.root
+
+    os.makedirs(opt.root, exist_ok=True)
+
+    opt.field = opt.field.split(',')
+    
+    # get file list
+    if os.path.exists(os.path.join(opt.root, 'metadata.csv')):
+        print('Loading previous metadata...')
+        metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv'))
+    else:
+        metadata = dataset_utils.get_metadata(**opt)
+        metadata.to_csv(os.path.join(opt.root, 'metadata.csv'), index=False)
+    
+    # merge downloaded
+    if opt.from_file:
+        downloaded_metadata = build_downloaded_metadata_from_files(
+            os.path.join(opt.download_root, 'raw'), metadata)
+    else:
+        downloaded_metadata = update_metadata(os.path.join(opt.download_root, 'raw'), opt)
+
+    # merge thumbnails
+    thumbnail_metadata = update_metadata(os.path.join(opt.thumbnail_root, 'thumbnails'), opt)
+    
+    # merge aesthetic scores
+    aesthetic_score_metadata = update_metadata(os.path.join(opt.root, 'aesthetic_scores'), opt)
+    
+    # merge render conditions
+    render_cond_metadata = update_metadata(os.path.join(opt.render_cond_root, 'renders_cond'), opt)
+
+    # merge mesh dumped
+    mesh_dumped_metadata = update_metadata(os.path.join(opt.mesh_dump_root, 'mesh_dumps'), opt)
+        
+    # merge pbr dumped
+    pbr_dumped_metadata = update_metadata(os.path.join(opt.pbr_dump_root, 'pbr_dumps'), opt)
+    
+    # merge asset stats
+    asset_stats_metadata = update_metadata(os.path.join(opt.root, 'asset_stats'), opt)
+        
+    # merge dual grid (original, no view transform)
+    dual_grid_resolutions = []
+    for dir in os.listdir(opt.dual_grid_root):
+        if os.path.isdir(os.path.join(opt.dual_grid_root, dir)) and dir.startswith('dual_grid_') and not dir.startswith('dual_grid_view_'):
+            dual_grid_resolutions.append(int(dir.split('_')[-1]))
+    dual_grid_metadata = {}
+    for res in dual_grid_resolutions:
+        dual_grid_metadata[res] = update_metadata(os.path.join(opt.dual_grid_root, f'dual_grid_{res}'), opt)
+    
+    # merge dual grid view (multi-view)
+    dual_grid_view_resolutions = []
+    for dir in os.listdir(opt.dual_grid_root):
+        if os.path.isdir(os.path.join(opt.dual_grid_root, dir)) and dir.startswith('dual_grid_view_'):
+            dual_grid_view_resolutions.append(int(dir.split('_')[-1]))
+    dual_grid_view_metadata = {}
+    for res in dual_grid_view_resolutions:
+        dual_grid_view_metadata[res] = update_metadata(os.path.join(opt.dual_grid_root, f'dual_grid_view_{res}'), opt)
+    
+    # merge pbr voxelized (single view)
+    pbr_voxel_resolutions = []
+    for dir in os.listdir(opt.pbr_voxel_root):
+        if os.path.isdir(os.path.join(opt.pbr_voxel_root, dir)) and dir.startswith('pbr_voxels_') and not dir.startswith('pbr_voxels_view_'):
+            pbr_voxel_resolutions.append(int(dir.split('_')[-1]))
+    pbr_voxel_metadata = {}
+    for res in pbr_voxel_resolutions:
+        pbr_voxel_metadata[res] = update_metadata(os.path.join(opt.pbr_voxel_root, f'pbr_voxels_{res}'), opt)
+    
+    # merge pbr voxelized view (multi-view)
+    # Supports both pbr_voxels_view_{res} and pbr_voxels_view_fix_{res} directory names
+    pbr_voxel_view_dirs = {}  # res -> dir_name
+    for dir in os.listdir(opt.pbr_voxel_root):
+        if os.path.isdir(os.path.join(opt.pbr_voxel_root, dir)) and dir.startswith('pbr_voxels_view_') and not dir.startswith('pbr_voxels_view_fix_'):
+            res = int(dir.split('_')[-1])
+            pbr_voxel_view_dirs[res] = dir
+        elif os.path.isdir(os.path.join(opt.pbr_voxel_root, dir)) and dir.startswith('pbr_voxels_view_fix_'):
+            res = int(dir.split('_')[-1])
+            pbr_voxel_view_dirs[res] = dir
+    pbr_voxel_view_resolutions = sorted(pbr_voxel_view_dirs.keys())
+    pbr_voxel_view_metadata = {}
+    for res in pbr_voxel_view_resolutions:
+        pbr_voxel_view_metadata[res] = update_metadata(os.path.join(opt.pbr_voxel_root, pbr_voxel_view_dirs[res]), opt)
+        
+    # merge ss latents
+    ss_latent_models = []
+    if os.path.exists(os.path.join(opt.ss_latent_root, 'ss_latents')):
+        ss_latent_models = os.listdir(os.path.join(opt.ss_latent_root, 'ss_latents'))
+    ss_latent_metadata = {}
+    for model in ss_latent_models:
+        ss_latent_metadata[model] = update_metadata(os.path.join(opt.ss_latent_root, f'ss_latents/{model}'), opt)
+        
+    # merge shape latents (original, no view transform)
+    shape_latent_models = []
+    if os.path.exists(os.path.join(opt.shape_latent_root, 'shape_latents')):
+        for dir in os.listdir(os.path.join(opt.shape_latent_root, 'shape_latents')):
+            if os.path.isdir(os.path.join(opt.shape_latent_root, 'shape_latents', dir)) and not _is_view_dir(dir):
+                shape_latent_models.append(dir)
+    shape_latent_metadata = {}
+    for model in shape_latent_models:
+        shape_latent_metadata[model] = update_metadata(os.path.join(opt.shape_latent_root, f'shape_latents/{model}'), opt)
+    
+    # merge shape latents view (multi-view, including _view and _view_fix)
+    shape_latent_view_models = []
+    if os.path.exists(os.path.join(opt.shape_latent_root, 'shape_latents')):
+        for dir in os.listdir(os.path.join(opt.shape_latent_root, 'shape_latents')):
+            if os.path.isdir(os.path.join(opt.shape_latent_root, 'shape_latents', dir)) and _is_view_dir(dir):
+                shape_latent_view_models.append(dir)
+    shape_latent_view_metadata = {}
+    for model in shape_latent_view_models:
+        shape_latent_view_metadata[model] = update_metadata(os.path.join(opt.shape_latent_root, f'shape_latents/{model}'), opt)
+        
+    # merge pbr latents (single view)
+    pbr_latent_models = []
+    if os.path.exists(os.path.join(opt.pbr_latent_root, 'pbr_latents')):
+        for dir in os.listdir(os.path.join(opt.pbr_latent_root, 'pbr_latents')):
+            if os.path.isdir(os.path.join(opt.pbr_latent_root, 'pbr_latents', dir)) and not _is_view_dir(dir):
+                pbr_latent_models.append(dir)
+    pbr_latent_metadata = {}
+    for model in pbr_latent_models:
+        pbr_latent_metadata[model] = update_metadata(os.path.join(opt.pbr_latent_root, f'pbr_latents/{model}'), opt)
+    
+    # merge pbr latents view (multi-view, including _view and _view_fix)
+    pbr_latent_view_models = []
+    if os.path.exists(os.path.join(opt.pbr_latent_root, 'pbr_latents')):
+        for dir in os.listdir(os.path.join(opt.pbr_latent_root, 'pbr_latents')):
+            if os.path.isdir(os.path.join(opt.pbr_latent_root, 'pbr_latents', dir)) and _is_view_dir(dir):
+                pbr_latent_view_models.append(dir)
+    pbr_latent_view_metadata = {}
+    for model in pbr_latent_view_models:
+        pbr_latent_view_metadata[model] = update_metadata(os.path.join(opt.pbr_latent_root, f'pbr_latents/{model}'), opt)
+
+    # Merge all sub-metadata back into main metadata and save
+    metadata = metadata.set_index('sha256')
+    sub_metadata_list = [
+        downloaded_metadata,
+        thumbnail_metadata,
+        aesthetic_score_metadata,
+        render_cond_metadata,
+        mesh_dumped_metadata,
+        pbr_dumped_metadata,
+        asset_stats_metadata,
+    ]
+    for res in dual_grid_resolutions:
+        sub_metadata_list.append(dual_grid_metadata.get(res))
+    for res in dual_grid_view_resolutions:
+        sub_metadata_list.append(dual_grid_view_metadata.get(res))
+    for res in pbr_voxel_resolutions:
+        sub_metadata_list.append(pbr_voxel_metadata.get(res))
+    for res in pbr_voxel_view_resolutions:
+        sub_metadata_list.append(pbr_voxel_view_metadata.get(res))
+    for model in ss_latent_models:
+        sub_metadata_list.append(ss_latent_metadata.get(model))
+    for model in shape_latent_models:
+        sub_metadata_list.append(shape_latent_metadata.get(model))
+    for model in shape_latent_view_models:
+        sub_metadata_list.append(shape_latent_view_metadata.get(model))
+    for model in pbr_latent_models:
+        sub_metadata_list.append(pbr_latent_metadata.get(model))
+    for model in pbr_latent_view_models:
+        sub_metadata_list.append(pbr_latent_view_metadata.get(model))
+    if metadata.index.duplicated().any():
+        metadata = metadata.groupby(level=0).first()
+    for sub in sub_metadata_list:
+        if sub is not None:
+            if 'sha256' in sub.columns:
+                sub = sub.set_index('sha256')
+            if sub.index.duplicated().any():
+                sub = sub.groupby(level=0).first()
+            metadata = metadata.combine_first(sub)
+    metadata = metadata.reset_index()
+    metadata.to_csv(os.path.join(opt.root, 'metadata.csv'), index=False)
+    print(f'Saved merged metadata with {len(metadata)} entries and columns: {list(metadata.columns)}')
+
+    # statistics
+    num_downloaded = downloaded_metadata['local_path'].count() if downloaded_metadata is not None else 0
+    with open(os.path.join(opt.root, 'statistics.txt'), 'w') as f:
+        f.write('Statistics:\n')
+        f.write(f'  - Number of assets: {len(metadata)}\n')
+        f.write(f'  - Number of assets downloaded: {num_downloaded}\n')
+        if thumbnail_metadata is not None:
+            f.write(f'  - Number of assets with thumbnails: {thumbnail_metadata["thumbnailed"].sum()}\n')
+        if aesthetic_score_metadata is not None:
+            f.write(f'  - Number of assets with aesthetic scores: {aesthetic_score_metadata["aesthetic_score"].count()}\n')
+        if render_cond_metadata is not None:
+            f.write(f'  - Number of assets with render conditions: {render_cond_metadata["cond_rendered"].count()}\n')
+        if mesh_dumped_metadata is not None:
+            f.write(f'  - Number of assets with mesh dumped: {mesh_dumped_metadata["mesh_dumped"].sum()}\n')
+        if pbr_dumped_metadata is not None:
+            f.write(f'  - Number of assets with PBR dumped: {pbr_dumped_metadata["pbr_dumped"].sum()}\n')
+        if asset_stats_metadata is not None:
+            f.write(f'  - Number of assets with asset stats: {len(asset_stats_metadata)}\n')
+        if len(dual_grid_resolutions) != 0:
+            f.write(f'  - Number of assets with dual grid:\n')
+            for res in dual_grid_resolutions:
+                if dual_grid_metadata[res] is not None:
+                    f.write(f'    - {res}: {dual_grid_metadata[res]["dual_grid_converted"].sum()}\n')
+        if len(dual_grid_view_resolutions) != 0:
+            f.write(f'  - Number of assets with dual grid view:\n')
+            for res in dual_grid_view_resolutions:
+                if dual_grid_view_metadata[res] is not None:
+                    col_name = f'dual_grid_view00_converted_{res}'
+                    if col_name in dual_grid_view_metadata[res].columns:
+                        f.write(f'    - {res}: {dual_grid_view_metadata[res][col_name].sum()}\n')
+                    else:
+                        f.write(f'    - {res}: {len(dual_grid_view_metadata[res])}\n')
+        if len(pbr_voxel_resolutions) != 0:
+            f.write(f'  - Number of assets with PBR voxelization:\n')
+            for res in sorted(pbr_voxel_resolutions):
+                if pbr_voxel_metadata[res] is not None:
+                    f.write(f'    - {res}: {pbr_voxel_metadata[res]["pbr_voxelized"].sum()}\n')
+        if len(pbr_voxel_view_resolutions) != 0:
+            f.write(f'  - Number of assets with PBR voxelization view:\n')
+            for res in sorted(pbr_voxel_view_resolutions):
+                if pbr_voxel_view_metadata[res] is not None:
+                    dir_name = pbr_voxel_view_dirs[res]
+                    col_name_old = 'pbr_voxelized_view00'
+                    col_name_new = f'pbr_voxelized_view_fix00_{res}'
+                    if col_name_old in pbr_voxel_view_metadata[res].columns:
+                        f.write(f'    - {dir_name}: {pbr_voxel_view_metadata[res][col_name_old].sum()}\n')
+                    elif col_name_new in pbr_voxel_view_metadata[res].columns:
+                        f.write(f'    - {dir_name}: {pbr_voxel_view_metadata[res][col_name_new].sum()}\n')
+                    else:
+                        f.write(f'    - {dir_name}: {len(pbr_voxel_view_metadata[res])}\n')
+        if len(ss_latent_models) != 0:
+            f.write(f'  - Number of assets with sparse structure latents:\n')
+            for model in ss_latent_models:
+                if ss_latent_metadata[model] is not None:
+                    if 'ss_latent_encoded' in ss_latent_metadata[model].columns:
+                        f.write(f'    - {model}: {ss_latent_metadata[model]["ss_latent_encoded"].sum()}\n')
+                    elif 'ss_latent_view00_encoded' in ss_latent_metadata[model].columns:
+                        f.write(f'    - {model}: {ss_latent_metadata[model]["ss_latent_view00_encoded"].sum()}\n')
+                    else:
+                        f.write(f'    - {model}: {len(ss_latent_metadata[model])}\n')
+        if len(shape_latent_models) != 0:
+            f.write(f'  - Number of assets with shape latents:\n')
+            for model in shape_latent_models:
+                if shape_latent_metadata[model] is not None:
+                    f.write(f'    - {model}: {shape_latent_metadata[model]["shape_latent_encoded"].sum()}\n')
+        if len(shape_latent_view_models) != 0:
+            f.write(f'  - Number of assets with shape latents view:\n')
+            for model in shape_latent_view_models:
+                if shape_latent_view_metadata[model] is not None:
+                    col_name = 'shape_latent_view00_encoded'
+                    if col_name in shape_latent_view_metadata[model].columns:
+                        f.write(f'    - {model}: {shape_latent_view_metadata[model][col_name].sum()}\n')
+                    else:
+                        f.write(f'    - {model}: {len(shape_latent_view_metadata[model])}\n')
+        if len(pbr_latent_models) != 0:
+            f.write(f'  - Number of assets with PBR latents:\n')
+            for model in pbr_latent_models:
+                if pbr_latent_metadata[model] is not None:
+                    f.write(f'    - {model}: {pbr_latent_metadata[model]["pbr_latent_encoded"].sum()}\n')
+        if len(pbr_latent_view_models) != 0:
+            f.write(f'  - Number of assets with PBR latents view:\n')
+            for model in pbr_latent_view_models:
+                if pbr_latent_view_metadata[model] is not None:
+                    col_name = 'pbr_latent_view00_encoded'
+                    if col_name in pbr_latent_view_metadata[model].columns:
+                        f.write(f'    - {model}: {pbr_latent_view_metadata[model][col_name].sum()}\n')
+                    else:
+                        f.write(f'    - {model}: {len(pbr_latent_view_metadata[model])}\n')
+        
+    with open(os.path.join(opt.root, 'statistics.txt'), 'r') as f:
+        print(f.read())
diff --git a/data_toolkit/build_texverse_metadata.py b/data_toolkit/build_texverse_metadata.py
new file mode 100644
index 0000000..366d548
--- /dev/null
+++ b/data_toolkit/build_texverse_metadata.py
@@ -0,0 +1,102 @@
+"""
+Build metadata.csv for TexVerse dataset from metadata.json.
+
+Output format matches ABO metadata.csv:
+    sha256, file_identifier, local_path
+
+Usage:
+    python data_toolkit/build_texverse_metadata.py \
+        --metadata_json /path/to/TexVerse/metadata.json \
+        --data_root /path/to/TexVerse \
+        --output /path/to/TexVerse/metadata.csv \
+        --max_workers 32
+"""
+
+import os
+import json
+import argparse
+import hashlib
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from tqdm import tqdm
+import pandas as pd
+
+
+def get_file_hash(file: str) -> str:
+    sha256 = hashlib.sha256()
+    with open(file, "rb") as f:
+        for byte_block in iter(lambda: f.read(4096), b""):
+            sha256.update(byte_block)
+    return sha256.hexdigest()
+
+
+def process_one(uid, glb_paths, data_root):
+    """Find the first existing GLB file for this uid and compute its sha256."""
+    for rel_path in glb_paths:
+        full_path = os.path.join(data_root, rel_path)
+        if os.path.exists(full_path):
+            sha256 = get_file_hash(full_path)
+            return {
+                'sha256': sha256,
+                'file_identifier': uid,
+                'local_path': rel_path,
+            }
+    return None
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Build metadata.csv for TexVerse')
+    parser.add_argument('--metadata_json', type=str, required=True,
+                        help='Path to TexVerse metadata.json')
+    parser.add_argument('--data_root', type=str, required=True,
+                        help='Root directory of TexVerse dataset (where glbs/ is)')
+    parser.add_argument('--output', type=str, required=True,
+                        help='Output path for metadata.csv')
+    parser.add_argument('--max_workers', type=int, default=16,
+                        help='Number of parallel workers')
+    args = parser.parse_args()
+
+    with open(args.metadata_json, 'r') as f:
+        metadata = json.load(f)
+
+    print(f'Total entries in metadata.json: {len(metadata)}')
+
+    # Load existing metadata.csv and skip already processed entries
+    existing_uids = set()
+    existing_records = []
+    if os.path.exists(args.output):
+        existing_df = pd.read_csv(args.output)
+        existing_uids = set(existing_df['file_identifier'].values)
+        existing_records = existing_df.to_dict('records')
+        print(f'Found existing metadata.csv with {len(existing_uids)} entries, skipping them')
+
+    # Filter out already processed uids
+    to_process = {uid: info for uid, info in metadata.items() if uid not in existing_uids}
+    print(f'New entries to process: {len(to_process)}')
+
+    if len(to_process) == 0:
+        print('Nothing to do, all entries already exist.')
+        return
+
+    new_records = []
+    with ProcessPoolExecutor(max_workers=args.max_workers) as executor:
+        futures = {
+            executor.submit(process_one, uid, info['glb_paths'], args.data_root): uid
+            for uid, info in to_process.items()
+        }
+        for future in tqdm(as_completed(futures), total=len(futures), desc='Building metadata'):
+            try:
+                result = future.result()
+                if result is not None:
+                    new_records.append(result)
+            except Exception as e:
+                uid = futures[future]
+                print(f'Error processing {uid}: {e}')
+
+    all_records = existing_records + new_records
+    df = pd.DataFrame.from_records(all_records, columns=['sha256', 'file_identifier', 'local_path'])
+    df.to_csv(args.output, index=False)
+    print(f'Added {len(new_records)} new entries, total {len(df)} entries in {args.output}')
+
+
+if __name__ == '__main__':
+    main()
diff --git a/data_toolkit/datasets/ABO.py b/data_toolkit/datasets/ABO.py
new file mode 100644
index 0000000..f783d83
--- /dev/null
+++ b/data_toolkit/datasets/ABO.py
@@ -0,0 +1,120 @@
+import os
+import re
+import argparse
+import tarfile
+from concurrent.futures import ThreadPoolExecutor
+from tqdm import tqdm
+import pandas as pd
+from utils import get_file_hash
+
+
+def add_args(parser: argparse.ArgumentParser):
+    pass
+
+
+def get_metadata(**kwargs):
+    metadata = pd.read_csv("hf://datasets/JeffreyXiang/TRELLIS-500K/ABO.csv")
+    return metadata
+        
+
+def download(metadata, root, **kwargs):    
+    output_dir  = root
+    os.makedirs(os.path.join(output_dir, 'raw'), exist_ok=True)
+
+    if not os.path.exists(os.path.join(output_dir, 'raw', 'abo-3dmodels.tar')):
+        try:
+            os.makedirs(os.path.join(output_dir, 'raw'), exist_ok=True)
+            os.system(f"wget -O {output_dir}/raw/abo-3dmodels.tar https://amazon-berkeley-objects.s3.amazonaws.com/archives/abo-3dmodels.tar")
+        except:
+            print("\033[93m")
+            print("Error downloading ABO dataset. Please check your internet connection and try again.")
+            print("Or, you can manually download the abo-3dmodels.tar file and place it in the {output_dir}/raw directory")
+            print("Visit https://amazon-berkeley-objects.s3.amazonaws.com/index.html for more information")
+            print("\033[0m")
+            raise FileNotFoundError("Error downloading ABO dataset")
+    
+    downloaded = {}
+    metadata = metadata.set_index("file_identifier")
+    with tarfile.open(os.path.join(output_dir, 'raw', 'abo-3dmodels.tar')) as tar:
+        with ThreadPoolExecutor(max_workers=1) as executor, \
+            tqdm(total=len(metadata), desc="Extracting") as pbar:
+            def worker(instance: str) -> str:
+                try:
+                    tar.extract(f"3dmodels/original/{instance}", path=os.path.join(output_dir, 'raw'))
+                    sha256 = get_file_hash(os.path.join(output_dir, 'raw/3dmodels/original', instance))
+                    pbar.update()
+                    return sha256
+                except Exception as e:
+                    pbar.update()
+                    print(f"Error extracting for {instance}: {e}")
+                    return None
+                
+            sha256s = executor.map(worker, metadata.index)
+            executor.shutdown(wait=True)
+
+    for k, sha256 in zip(metadata.index, sha256s):
+        if sha256 is not None:
+            if sha256 == metadata.loc[k, "sha256"]:
+                downloaded[sha256] = os.path.join('raw/3dmodels/original', k)
+            else:
+                print(f"Error downloading {k}: sha256s do not match")
+
+    return pd.DataFrame(downloaded.items(), columns=['sha256', 'local_path'])
+
+
+def _process_instance(args):
+    """Worker function for ProcessPoolExecutor (must be top-level for pickling)"""
+    import os
+    metadatum, output_dir, func = args
+    try:
+        local_path = metadatum['local_path']
+        sha256 = metadatum['sha256']
+        file = os.path.join(output_dir, local_path)
+        record = func(file, sha256)
+        return record
+    except Exception as e:
+        print(f"Error processing object {metadatum.get('sha256', '?')}: {e}")
+        return None
+
+
+def foreach_instance(metadata, output_dir, func, max_workers=None, desc='Processing objects') -> pd.DataFrame:
+    import os
+    from concurrent.futures import ProcessPoolExecutor, as_completed
+    from tqdm import tqdm
+    
+    # load metadata
+    metadata = metadata.to_dict('records')
+
+    max_workers = max_workers or os.cpu_count()
+    records = []
+    
+    # Track processed/skipped counts
+    total_processed = 0
+    total_skipped = 0
+    
+    try:
+        with ProcessPoolExecutor(max_workers=max_workers) as executor:
+            futures = {
+                executor.submit(_process_instance, (m, output_dir, func)): m['sha256']
+                for m in metadata
+            }
+            pbar = tqdm(as_completed(futures), total=len(futures), desc=desc)
+            for future in pbar:
+                try:
+                    r = future.result()
+                    if r is not None:
+                        records.append(r)
+                        # Update stats
+                        if '_processed_count' in r:
+                            total_processed += r['_processed_count']
+                        if '_skipped_count' in r:
+                            total_skipped += r['_skipped_count']
+                        # Update progress bar display
+                        pbar.set_postfix(processed=total_processed, skipped=total_skipped, refresh=False)
+                except Exception as e:
+                    sha256 = futures[future]
+                    print(f"Error processing object {sha256}: {e}")
+    except Exception as e:
+        print(f"Error happened during processing: {e}")
+        
+    return pd.DataFrame.from_records(records)
\ No newline at end of file
diff --git a/data_toolkit/datasets/ObjaverseXL.py b/data_toolkit/datasets/ObjaverseXL.py
new file mode 100644
index 0000000..6320912
--- /dev/null
+++ b/data_toolkit/datasets/ObjaverseXL.py
@@ -0,0 +1,133 @@
+import os
+import argparse
+from concurrent.futures import ThreadPoolExecutor
+from tqdm import tqdm
+import pandas as pd
+import objaverse.xl as oxl
+from utils import get_file_hash
+
+
+def add_args(parser: argparse.ArgumentParser):
+    parser.add_argument('--source', type=str, default='sketchfab',
+                        help='Data source to download annotations from (github, sketchfab)')
+
+
+def get_metadata(source, **kwargs):
+    if source == 'sketchfab':
+        metadata = pd.read_csv("hf://datasets/JeffreyXiang/TRELLIS-500K/ObjaverseXL_sketchfab.csv")
+    elif source == 'github':
+        metadata = pd.read_csv("hf://datasets/JeffreyXiang/TRELLIS-500K/ObjaverseXL_github.csv")
+    else:
+        raise ValueError(f"Invalid source: {source}")
+    return metadata
+        
+
+def download(metadata, output_dir, **kwargs):    
+    os.makedirs(os.path.join(output_dir, 'raw'), exist_ok=True)
+
+    # download annotations
+    annotations = oxl.get_annotations()
+    annotations = annotations[annotations['sha256'].isin(metadata['sha256'].values)]
+    
+    # download and render objects
+    file_paths = oxl.download_objects(
+        annotations,
+        download_dir=os.path.join(output_dir, "raw"),
+        save_repo_format="zip",
+    )
+    
+    downloaded = {}
+    metadata = metadata.set_index("file_identifier")
+    for k, v in file_paths.items():
+        sha256 = metadata.loc[k, "sha256"]
+        downloaded[sha256] = os.path.relpath(v, output_dir)
+
+    return pd.DataFrame(downloaded.items(), columns=['sha256', 'local_path'])
+
+
+def _process_instance(args):
+    """Worker function for ProcessPoolExecutor (must be top-level for pickling)"""
+    import os, tempfile, zipfile
+    metadatum, output_dir, func = args
+    try:
+        local_path = metadatum['local_path']
+        sha256 = metadatum['sha256']
+        
+        direct_file_path = os.path.join(output_dir, local_path)
+        if os.path.exists(direct_file_path):
+            file = direct_file_path
+            record = func(file, sha256)
+        elif local_path.startswith('raw/github/repos/'):
+            path_parts = local_path.split('/')
+            file_name = os.path.join(*path_parts[5:])
+            zip_file = os.path.join(output_dir, *path_parts[:5])
+            if os.path.exists(zip_file):
+                with tempfile.TemporaryDirectory() as tmp_dir:
+                    with zipfile.ZipFile(zip_file, 'r') as zip_ref:
+                        zip_ref.extractall(tmp_dir)
+                    file = os.path.join(tmp_dir, file_name)
+                    record = func(file, sha256)
+            else:
+                # zip file not found, pass local_path directly (for tasks like dual_grid_view that don't need the original file)
+                file = local_path
+                record = func(file, sha256)
+        else:
+            file = os.path.join(output_dir, local_path)
+            record = func(file, sha256)
+        return record
+    except Exception as e:
+        print(f"Error processing object {metadatum.get('sha256', '?')}: {e}")
+        return None
+
+
+def foreach_instance(metadata, output_dir, func, max_workers=None, desc='Processing objects', log_interval=500, timeout=None) -> pd.DataFrame:
+    print("================")
+    import os
+    from concurrent.futures import ProcessPoolExecutor, as_completed, TimeoutError
+    from tqdm import tqdm
+    
+    # load metadata
+    metadata = metadata.to_dict('records')
+
+    max_workers = max_workers or os.cpu_count()
+    records = []
+    
+    # Track processed/skipped counts
+    total_processed = 0
+    total_skipped = 0
+    timeout_count = 0
+    
+    try:
+        with ProcessPoolExecutor(max_workers=max_workers) as executor:
+            futures = {
+                executor.submit(_process_instance, (m, output_dir, func)): m['sha256']
+                for m in metadata
+            }
+            pbar = tqdm(as_completed(futures), total=len(futures), desc=desc)
+            for future in pbar:
+                sha256 = futures[future]
+                try:
+                    r = future.result(timeout=timeout)
+                    if r is not None:
+                        records.append(r)
+                        # Update stats
+                        if '_processed_count' in r:
+                            total_processed += r['_processed_count']
+                        if '_skipped_count' in r:
+                            total_skipped += r['_skipped_count']
+                        # Update progress bar display
+                        pbar.set_postfix(processed=total_processed, skipped=total_skipped, timeout=timeout_count, refresh=False)
+                except TimeoutError:
+                    timeout_count += 1
+                    print(f"Timeout processing object {sha256} (>{timeout}s)")
+                    records.append({'sha256': sha256, 'error': f'Timeout (>{timeout}s)'})
+                    pbar.set_postfix(processed=total_processed, skipped=total_skipped, timeout=timeout_count, refresh=False)
+                except Exception as e:
+                    print(f"Error processing object {sha256}: {e}")
+    except Exception as e:
+        print(f"Error happened during processing: {e}")
+    
+    if timeout_count > 0:
+        print(f"Total timeout: {timeout_count} objects")
+        
+    return pd.DataFrame.from_records(records)
diff --git a/data_toolkit/datasets/TexVerse.py b/data_toolkit/datasets/TexVerse.py
new file mode 100644
index 0000000..8600ce9
--- /dev/null
+++ b/data_toolkit/datasets/TexVerse.py
@@ -0,0 +1,69 @@
+import os
+import argparse
+from concurrent.futures import ProcessPoolExecutor, as_completed
+from tqdm import tqdm
+import pandas as pd
+
+
+def add_args(parser: argparse.ArgumentParser):
+    pass
+
+
+def get_metadata(**kwargs):
+    raise NotImplementedError("TexVerse metadata should be built from metadata.json using build_texverse_metadata.py")
+
+
+def download(metadata, root, **kwargs):
+    raise NotImplementedError("TexVerse GLB files are already available locally. No download needed.")
+
+
+def _process_instance(args):
+    """Worker function for ProcessPoolExecutor (must be top-level for pickling)"""
+    metadatum, output_dir, func = args
+    try:
+        local_path = metadatum['local_path']
+        sha256 = metadatum['sha256']
+        file = os.path.join(output_dir, local_path)
+        record = func(file, sha256)
+        return record
+    except Exception as e:
+        print(f"Error processing object {metadatum.get('sha256', '?')}: {e}")
+        return None
+
+
+def foreach_instance(metadata, output_dir, func, max_workers=None, desc='Processing objects', timeout=None) -> pd.DataFrame:
+    import os
+    from concurrent.futures import ProcessPoolExecutor, as_completed, TimeoutError
+    from tqdm import tqdm
+
+    metadata = metadata.to_dict('records')
+
+    max_workers = max_workers or os.cpu_count()
+    records = []
+    timeout_count = 0
+
+    try:
+        with ProcessPoolExecutor(max_workers=max_workers) as executor:
+            futures = {
+                executor.submit(_process_instance, (m, output_dir, func)): m['sha256']
+                for m in metadata
+            }
+            for future in tqdm(as_completed(futures), total=len(futures), desc=desc):
+                sha256 = futures[future]
+                try:
+                    r = future.result(timeout=timeout)
+                    if r is not None:
+                        records.append(r)
+                except TimeoutError:
+                    timeout_count += 1
+                    print(f"Timeout processing object {sha256} (>{timeout}s)")
+                    records.append({'sha256': sha256, 'error': f'Timeout (>{timeout}s)'})
+                except Exception as e:
+                    print(f"Error processing object {sha256}: {e}")
+    except Exception as e:
+        print(f"Error happened during processing: {e}")
+
+    if timeout_count > 0:
+        print(f"Total timeout: {timeout_count} objects")
+
+    return pd.DataFrame.from_records(records)
diff --git a/data_toolkit/download.py b/data_toolkit/download.py
new file mode 100644
index 0000000..5697681
--- /dev/null
+++ b/data_toolkit/download.py
@@ -0,0 +1,64 @@
+import os
+import copy
+import sys
+import importlib
+import argparse
+import pandas as pd
+from easydict import EasyDict as edict
+
+if __name__ == '__main__':
+    dataset_utils = importlib.import_module(f'datasets.{sys.argv[1]}')
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--root', type=str, required=True,
+                        help='Directory to save the metadata')
+    parser.add_argument('--download_root', type=str, default=None,
+                        help='Directory to download the objects')
+    parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
+                        help='Filter objects with aesthetic score lower than this value')
+    parser.add_argument('--check_only', action='store_true',
+                        help='Only check if the objects are already downloaded')
+    parser.add_argument('--instances', type=str, default=None,
+                        help='Instances to process')
+    dataset_utils.add_args(parser)
+    parser.add_argument('--rank', type=int, default=0)
+    parser.add_argument('--world_size', type=int, default=1)
+    opt = parser.parse_args(sys.argv[2:])
+    opt = edict(vars(opt))
+    opt.download_root = opt.download_root or opt.root
+
+    os.makedirs(opt.root, exist_ok=True)
+    os.makedirs(opt.download_root, exist_ok=True)
+    os.makedirs(os.path.join(opt.download_root, 'raw', 'new_records'), exist_ok=True)
+
+    # get file list
+    if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
+        raise ValueError('metadata.csv not found')
+    metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
+    if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256'))
+    if os.path.exists(os.path.join(opt.download_root, 'raw', 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.download_root, 'raw', 'metadata.csv')).set_index('sha256'))
+    metadata = metadata.reset_index()
+    if opt.instances is None:
+        if opt.filter_low_aesthetic_score is not None:
+            metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
+        if 'local_path' in metadata.columns:
+            metadata = metadata[metadata['local_path'].isna()]
+    else:
+        if os.path.exists(opt.instances):
+            with open(opt.instances, 'r') as f:
+                instances = f.read().splitlines()
+        else:
+            instances = opt.instances.split(',')
+        metadata = metadata[metadata['sha256'].isin(instances)]
+
+    start = len(metadata) * opt.rank // opt.world_size
+    end = len(metadata) * (opt.rank + 1) // opt.world_size
+    metadata = metadata[start:end]
+                
+    print(f'Processing {len(metadata)} objects...')
+
+    # process objects
+    downloaded = dataset_utils.download(metadata, output_dir=opt.download_root, **opt)
+    downloaded.to_csv(os.path.join(opt.download_root, 'raw', 'new_records', f'part_{opt.rank}.csv'), index=False)
diff --git a/data_toolkit/dual_grid.py b/data_toolkit/dual_grid.py
new file mode 100644
index 0000000..f9a3621
--- /dev/null
+++ b/data_toolkit/dual_grid.py
@@ -0,0 +1,170 @@
+import os
+import sys
+import importlib
+import argparse
+import pandas as pd
+import numpy as np
+import torch
+import pickle
+import o_voxel
+from easydict import EasyDict as edict
+from functools import partial
+
+
+def _dual_grid_mesh(file, metadatum, mesh_dump_root, root):
+    sha256 = metadatum['sha256']
+    try:
+        pack = {'sha256': sha256}
+        data = None
+        for res in opt.resolution:
+            need_process = False
+
+            # check if already processed
+            if os.path.exists(os.path.join(root, f'dual_grid_{res}', f'{sha256}.vxz')):
+                try:
+                    info = o_voxel.io.read_vxz_info(os.path.join(root, f'dual_grid_{res}', f'{sha256}.vxz'))
+                    pack[f'dual_grid_converted_{res}'] = True
+                    pack[f'dual_grid_size_{res}'] = info['num_voxel']
+                except Exception as e:
+                    print(f'Error reading {sha256}.vxz: {e}')
+                    need_process = True
+            else:
+                need_process = True
+
+            # process mesh
+            if need_process:
+                if data is None:
+                    with open(os.path.join(mesh_dump_root, 'mesh_dumps', f'{sha256}.pickle'), 'rb') as f:
+                        dump = pickle.load(f)
+                    start = 0
+                    vertices = []
+                    faces = []
+                    for obj in dump['objects']:
+                        if obj['vertices'].size == 0 or obj['faces'].size == 0:
+                            continue
+                        vertices.append(obj['vertices'])
+                        faces.append(obj['faces'] + start)
+                        start += len(obj['vertices'])
+                    vertices = torch.from_numpy(np.concatenate(vertices, axis=0)).float()
+                    faces = torch.from_numpy(np.concatenate(faces, axis=0)).long()
+                    vertices_min = vertices.min(dim=0)[0]
+                    vertices_max = vertices.max(dim=0)[0]
+                    center = (vertices_min + vertices_max) / 2
+                    scale = 0.99999 / (vertices_max - vertices_min).max()
+                    vertices = (vertices - center) * scale
+                    assert torch.all(vertices >= -0.5) and torch.all(vertices <= 0.5), 'vertices out of range'
+                    data = {'vertices': vertices, 'faces': faces}
+
+                voxel_indices, dual_vertices, intersected = o_voxel.convert.mesh_to_flexible_dual_grid(
+                    **data,
+                    grid_size=res,
+                    aabb=[[-0.5,-0.5,-0.5],[0.5,0.5,0.5]],
+                    face_weight=1.0,
+                    boundary_weight=0.2,
+                    regularization_weight=1e-2,
+                    timing=False,
+                )
+                dual_vertices = dual_vertices * res - voxel_indices
+                assert torch.all(dual_vertices >= -1e-3) and torch.all(dual_vertices <= 1+1e-3), 'dual_vertices out of range'
+                dual_vertices = torch.clamp(dual_vertices, 0, 1)
+                dual_vertices = (dual_vertices * 255).type(torch.uint8)
+                intersected = (intersected[:, 0:1] + 2 * intersected[:, 1:2] + 4 * intersected[:, 2:3]).type(torch.uint8)
+                
+                o_voxel.io.write_vxz(
+                    os.path.join(root, f'dual_grid_{res}', f'{sha256}.vxz'),
+                    voxel_indices,
+                    {'vertices': dual_vertices, 'intersected': intersected},
+                )
+
+                pack[f'dual_grid_converted_{res}'] = True
+                pack[f'dual_grid_size_{res}'] = len(dual_vertices)
+
+        return pack
+    except Exception as e:
+        print(f'Error processing {sha256}: {e}')
+        return {'sha256': sha256, 'error': str(e)}
+
+
+if __name__ == '__main__':
+    dataset_utils = importlib.import_module(f'datasets.{sys.argv[1]}')
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--root', type=str, required=True,
+                        help='Directory to save the metadata')
+    parser.add_argument('--mesh_dump_root', type=str, default=None,
+                        help='Directory to load mesh dumps')
+    parser.add_argument('--dual_grid_root', type=str, default=None,
+                        help='Directory to save dual grids')
+    parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
+                        help='Filter objects with aesthetic score lower than this value')
+    parser.add_argument('--instances', type=str, default=None,
+                        help='Instances to process')
+    dataset_utils.add_args(parser)
+    parser.add_argument('--rank', type=int, default=0)
+    parser.add_argument('--resolution', type=str, default=256)
+    parser.add_argument('--world_size', type=int, default=1)
+    parser.add_argument('--max_workers', type=int, default=0)
+    opt = parser.parse_args(sys.argv[2:])
+    opt = edict(vars(opt))
+    opt.resolution = [int(x) for x in opt.resolution.split(',')]
+    opt.mesh_dump_root = opt.mesh_dump_root or opt.root
+    opt.dual_grid_root = opt.dual_grid_root or opt.root
+
+    for res in opt.resolution:
+        os.makedirs(os.path.join(opt.dual_grid_root, f'dual_grid_{res}', 'new_records'), exist_ok=True)
+
+    # get file list
+    if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
+        raise ValueError('metadata.csv not found')
+    metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
+    if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256'))
+    if os.path.exists(os.path.join(opt.mesh_dump_root, 'mesh_dumps', 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.mesh_dump_root, 'mesh_dumps', 'metadata.csv')).set_index('sha256'))
+    for res in opt.resolution:
+        if os.path.exists(os.path.join(opt.dual_grid_root, f'dual_grid_{res}', 'metadata.csv')):
+            dual_grid_metadata = pd.read_csv(os.path.join(opt.dual_grid_root, f'dual_grid_{res}', 'metadata.csv')).set_index('sha256')
+            dual_grid_metadata = dual_grid_metadata.rename(columns={'dual_grid_converted': f'dual_grid_converted_{res}', 'dual_grid_size': f'dual_grid_size_{res}'})
+            metadata = metadata.combine_first(dual_grid_metadata)
+    metadata = metadata.reset_index()
+    if opt.instances is None:
+        if opt.filter_low_aesthetic_score is not None:
+            metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
+        metadata = metadata[metadata['mesh_dumped'] == True]
+        mask = np.zeros(len(metadata), dtype=bool)
+        for res in opt.resolution:
+            if f'dual_grid_converted_{res}' in metadata.columns:
+                mask |= metadata[f'dual_grid_converted_{res}'] != True
+            else:
+                mask[:] = True
+                break
+        metadata = metadata[mask]
+    else:
+        if os.path.exists(opt.instances):
+            with open(opt.instances, 'r') as f:
+                instances = f.read().splitlines()
+        else:
+            instances = opt.instances.split(',')
+        metadata = metadata[metadata['sha256'].isin(instances)]
+
+    start = len(metadata) * opt.rank // opt.world_size
+    end = len(metadata) * (opt.rank + 1) // opt.world_size
+    metadata = metadata[start:end]
+    
+    print(f'Processing {len(metadata)} objects...')
+
+    # process objects
+    func = partial(_dual_grid_mesh, root=opt.dual_grid_root, mesh_dump_root=opt.mesh_dump_root)
+    dual_grids = dataset_utils.foreach_instance(metadata, None, func, max_workers=opt.max_workers, no_file=True, desc='Dual griding')
+    if 'error' in dual_grids.columns:
+        errors = dual_grids[dual_grids[f'error'].notna()]
+        with open('errors.txt', 'w') as f:
+            f.write('\n'.join(errors['sha256'].tolist()))
+    for res in opt.resolution:
+        if f'dual_grid_converted_{res}' in dual_grids.columns:
+            dual_grid_metadata = dual_grids[dual_grids[f'dual_grid_converted_{res}'] == True]
+            if len(dual_grid_metadata) > 0:
+                dual_grid_metadata = dual_grid_metadata[['sha256', f'dual_grid_converted_{res}', f'dual_grid_size_{res}']]
+                dual_grid_metadata = dual_grid_metadata.rename(columns={f'dual_grid_converted_{res}': 'dual_grid_converted', f'dual_grid_size_{res}': 'dual_grid_size'})
+                dual_grid_metadata.to_csv(os.path.join(opt.dual_grid_root, f'dual_grid_{res}', 'new_records', f'part_{opt.rank}.csv'), index=False)
+    
\ No newline at end of file
diff --git a/data_toolkit/dual_grid_view.py b/data_toolkit/dual_grid_view.py
new file mode 100644
index 0000000..3830b4d
--- /dev/null
+++ b/data_toolkit/dual_grid_view.py
@@ -0,0 +1,350 @@
+"""
+dual_grid_view.py - Multi-view transform dual grid processing
+Extends dual_grid.py with scale and mesh rotation logic
+Based on test_ovoxel_transform.py implementation
+"""
+import os
+import sys
+import importlib
+import argparse
+import json
+import pandas as pd
+import numpy as np
+import torch
+import pickle
+import o_voxel
+from easydict import EasyDict as edict
+from functools import partial
+from utils import get_new_camera_matrix, transform_mesh, sphere_normalize_torch
+
+
+def _dual_grid_mesh_view(file, sha256, mesh_dump_root, transform_root, root, view_indices=None):
+    """
+    Process multi-view dual grid conversion for a single sha256.
+    
+    Args:
+        file: local_path from metadata
+        sha256: sha256 string
+        mesh_dump_root: directory containing mesh dump files
+        transform_root: directory containing transform json files
+        root: output directory for dual grids
+        view_indices: list of view indices to process, None for all views
+    """
+    try:
+        pack = {'sha256': sha256}
+        vertices_sphere = None
+        sphere_radius = None
+        faces = None
+        
+        # Load transforms
+        transform_path = os.path.join(transform_root, sha256, 'transforms.json')
+        if not os.path.exists(transform_path):
+            print(f'Transform file not found for {sha256}, skipping')
+            return {'sha256': sha256, 'error': 'Transform file not found'}
+        
+        with open(transform_path, 'r') as f:
+            transforms_json = json.load(f)
+        transform_mats = transforms_json['frames']
+        
+        # Determine views to process
+        if view_indices is None:
+            view_indices = list(range(len(transform_mats)))
+        else:
+            view_indices = [i for i in view_indices if i < len(transform_mats)]
+        
+        # Track processed and skipped counts
+        processed_count = 0
+        skipped_count = 0
+        
+        for view_idx in view_indices:
+            for res in opt.resolution:
+                need_process = False
+                
+                # Check if already processed
+                # Path structure: dual_grid_view_{res}/{sha256}/view{idx:02d}.vxz
+                sha256_dir = os.path.join(root, f'dual_grid_view_{res}', sha256)
+                vxz_path = os.path.join(sha256_dir, f'view{view_idx:02d}.vxz')
+                if os.path.exists(vxz_path):
+                    try:
+                        info = o_voxel.io.read_vxz_info(vxz_path)
+                        pack[f'dual_grid_view{view_idx:02d}_converted_{res}'] = True
+                        pack[f'dual_grid_view{view_idx:02d}_size_{res}'] = info['num_voxel']
+                        skipped_count += 1
+                    except Exception as e:
+                        print(f'Error reading {sha256}/view{view_idx:02d}.vxz: {e}')
+                        need_process = True
+                else:
+                    need_process = True
+                
+                # Process mesh
+                if need_process:
+                    # Lazy load mesh
+                    if vertices_sphere is None:
+                        mesh_file = os.path.join(mesh_dump_root, 'mesh_dumps', f'{sha256}.pickle')
+                        if not os.path.exists(mesh_file):
+                            print(f'Mesh dump not found for {sha256}, skipping')
+                            return {'sha256': sha256, 'error': 'Mesh dump not found'}
+                        
+                        with open(mesh_file, 'rb') as f:
+                            dump = pickle.load(f)
+                        
+                        start = 0
+                        vertices_list = []
+                        faces_list = []
+                        for obj in dump['objects']:
+                            if obj['vertices'].size == 0 or obj['faces'].size == 0:
+                                continue
+                            vertices_list.append(obj['vertices'])
+                            faces_list.append(obj['faces'] + start)
+                            start += len(obj['vertices'])
+                        
+                        if len(vertices_list) == 0:
+                            print(f'No valid mesh data for {sha256}, skipping')
+                            return {'sha256': sha256, 'error': 'No valid mesh data'}
+                        
+                        vertices = torch.from_numpy(np.concatenate(vertices_list, axis=0)).float().contiguous()
+                        faces = torch.from_numpy(np.concatenate(faces_list, axis=0)).long().contiguous()
+                        
+                        # Sphere normalization (for multi-view transform) - CPU only
+                        vertices_sphere, sphere_center, sphere_radius = sphere_normalize_torch(vertices)
+                    
+                    # Get transform for current view
+                    transform = transform_mats[view_idx]
+                    
+                    # Multi-view transform - CPU only
+                    transformed_vertices = transform_mesh(vertices_sphere, transform)
+                    
+                    # Post-transform normalization: scale by abs max to [-0.5, 0.5]^3
+                    # Only scale, no center shift, to preserve relative model position
+                    abs_max = transformed_vertices.abs().max().item()
+                    box_scale = 0.49999 / abs_max  # Normalize to [-0.5, 0.5] range
+                    transformed_normalized = transformed_vertices * box_scale
+                    transformed_normalized_cpu = transformed_normalized.contiguous()
+                    
+                    # Compute total scale (from original mesh to final normalized mesh)
+                    total_scale = box_scale / sphere_radius.item()
+                    
+                    # Validate range
+                    assert torch.all(transformed_normalized_cpu >= -0.5) and torch.all(transformed_normalized_cpu <= 0.5), \
+                        f'vertices out of range for {sha256} view {view_idx}'
+                    
+                    # Ensure vertices and faces are on CPU with correct types and contiguous memory
+                    # CPU only, consistent with process_dual_grid in test_ovoxel_transform.py
+                    vertices_for_grid = transformed_normalized_cpu.float().contiguous()
+                    faces_for_grid = faces.long().contiguous()
+                    data_for_grid = {'vertices': vertices_for_grid, 'faces': faces_for_grid}
+                    
+                    # Dual grid encoding
+                    voxel_indices, dual_vertices, intersected = o_voxel.convert.mesh_to_flexible_dual_grid(
+                        **data_for_grid,
+                        grid_size=res,
+                        aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
+                        face_weight=1.0,
+                        boundary_weight=0.2,
+                        regularization_weight=1e-2,
+                        timing=False,
+                    )
+                    
+                    # Convert to intra-voxel offsets and quantize
+                    dual_vertices = dual_vertices.float()
+                    voxel_indices_float = voxel_indices.float()
+                    dual_vertices = dual_vertices * res - voxel_indices_float
+                    assert torch.all(dual_vertices >= -1e-3) and torch.all(dual_vertices <= 1+1e-3), \
+                        f'dual_vertices out of range for {sha256} view {view_idx}'
+                    dual_vertices = torch.clamp(dual_vertices, 0, 1)
+                    dual_vertices = (dual_vertices * 255).type(torch.uint8)
+                    intersected = (intersected[:, 0:1] + 2 * intersected[:, 1:2] + 4 * intersected[:, 2:3]).type(torch.uint8)
+                    
+                    # Save .vxz file
+                    os.makedirs(sha256_dir, exist_ok=True)
+                    o_voxel.io.write_vxz(
+                        vxz_path,
+                        voxel_indices,
+                        {'vertices': dual_vertices, 'intersected': intersected},
+                    )
+                    
+                    # Save scale info
+                    scale_path = os.path.join(sha256_dir, f'view{view_idx:02d}_scale.json')
+                    scale_info = {
+                        'sha256': sha256,
+                        'view_idx': view_idx,
+                        'total_scale': total_scale,
+                        'sphere_radius': sphere_radius.item(),
+                        'box_scale': box_scale,
+                    }
+                    with open(scale_path, 'w') as f:
+                        json.dump(scale_info, f, indent=2)
+                    
+                    pack[f'dual_grid_view{view_idx:02d}_converted_{res}'] = True
+                    pack[f'dual_grid_view{view_idx:02d}_size_{res}'] = len(voxel_indices)
+                    pack[f'dual_grid_view{view_idx:02d}_scale_{res}'] = total_scale
+                    processed_count += 1
+        
+        # Record processing stats
+        pack['_processed_count'] = processed_count
+        pack['_skipped_count'] = skipped_count
+        
+        return pack
+    
+    except Exception as e:
+        print(f'Error processing {sha256}: {e}')
+        import traceback
+        traceback.print_exc()
+        return {'sha256': sha256, 'error': str(e)}
+
+
+if __name__ == '__main__':
+    dataset_utils = importlib.import_module(f'datasets.{sys.argv[1]}')
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--root', type=str, required=True,
+                        help='Directory to save the metadata')
+    parser.add_argument('--mesh_dump_root', type=str, default=None,
+                        help='Directory to load mesh dumps')
+    parser.add_argument('--transform_root', type=str, default=None,
+                        help='Directory to load transform json files (renders_cond)')
+    parser.add_argument('--dual_grid_root', type=str, default=None,
+                        help='Directory to save dual grids')
+    parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
+                        help='Filter objects with aesthetic score lower than this value')
+    parser.add_argument('--instances', type=str, default=None,
+                        help='Instances to process')
+    parser.add_argument('--view_indices', type=str, default=None,
+                        help='View indices to process, e.g., "0,1,2" or "0-5". None for all views')
+    dataset_utils.add_args(parser)
+    parser.add_argument('--rank', type=int, default=0)
+    parser.add_argument('--resolution', type=str, default='256')
+    parser.add_argument('--world_size', type=int, default=1)
+    parser.add_argument('--max_workers', type=int, default=0)
+    opt = parser.parse_args(sys.argv[2:])
+    opt = edict(vars(opt))
+    opt.resolution = [int(x) for x in opt.resolution.split(',')]
+    opt.mesh_dump_root = opt.mesh_dump_root or opt.root
+    opt.transform_root = opt.transform_root or os.path.join(opt.root, 'renders_cond')
+    opt.dual_grid_root = opt.dual_grid_root or opt.root
+    
+    # Parse view_indices
+    view_indices = None
+    if opt.view_indices is not None:
+        view_indices = []
+        for part in opt.view_indices.split(','):
+            if '-' in part:
+                start, end = map(int, part.split('-'))
+                view_indices.extend(range(start, end + 1))
+            else:
+                view_indices.append(int(part))
+        view_indices = list(set(view_indices))  # Deduplicate
+        view_indices.sort()
+
+    for res in opt.resolution:
+        os.makedirs(os.path.join(opt.dual_grid_root, f'dual_grid_view_{res}', 'new_records'), exist_ok=True)
+
+    # Get file list
+    if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
+        raise ValueError('metadata.csv not found')
+    metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
+    if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256'))
+    if os.path.exists(os.path.join(opt.mesh_dump_root, 'mesh_dumps', 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.mesh_dump_root, 'mesh_dumps', 'metadata.csv')).set_index('sha256'))
+    
+    # Check already processed dual_grid_view
+    for res in opt.resolution:
+        if os.path.exists(os.path.join(opt.dual_grid_root, f'dual_grid_view_{res}', 'metadata.csv')):
+            dual_grid_metadata = pd.read_csv(os.path.join(opt.dual_grid_root, f'dual_grid_view_{res}', 'metadata.csv')).set_index('sha256')
+            metadata = metadata.combine_first(dual_grid_metadata)
+    
+    metadata = metadata.reset_index()
+    
+    if opt.instances is None:
+        if opt.filter_low_aesthetic_score is not None:
+            metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
+        metadata = metadata[metadata['mesh_dumped'] == True]
+        
+        # Filter out objects with all views already processed
+        if view_indices is not None:
+            for res in opt.resolution:
+                # Check if each specified view is already processed
+                all_views_done_col = f'_all_views_done_{res}'
+                metadata[all_views_done_col] = True
+                for view_idx in view_indices:
+                    col_name = f'dual_grid_view{view_idx:02d}_converted_{res}'
+                    if col_name in metadata.columns:
+                        metadata[all_views_done_col] = metadata[all_views_done_col] & (metadata[col_name] == True)
+                    else:
+                        metadata[all_views_done_col] = False
+                        break
+            
+            # Keep objects with at least one incomplete resolution
+            any_incomplete = None
+            for res in opt.resolution:
+                all_views_done_col = f'_all_views_done_{res}'
+                if all_views_done_col in metadata.columns:
+                    if any_incomplete is None:
+                        any_incomplete = ~metadata[all_views_done_col]
+                    else:
+                        any_incomplete = any_incomplete | ~metadata[all_views_done_col]
+            
+            if any_incomplete is not None:
+                before_filter = len(metadata)
+                metadata = metadata[any_incomplete]
+                print(f'Filtered out {before_filter - len(metadata)} already completed objects')
+    else:
+        if os.path.exists(opt.instances):
+            with open(opt.instances, 'r') as f:
+                instances = f.read().splitlines()
+        else:
+            instances = opt.instances.split(',')
+        metadata = metadata[metadata['sha256'].isin(instances)]
+
+    start = len(metadata) * opt.rank // opt.world_size
+    end = len(metadata) * (opt.rank + 1) // opt.world_size
+    metadata = metadata[start:end]
+    
+    print(f'Processing {len(metadata)} objects...')
+    if view_indices:
+        print(f'View indices to process: {view_indices}')
+    else:
+        print('Processing all available views')
+
+    # Process objects
+    func = partial(_dual_grid_mesh_view, 
+                   root=opt.dual_grid_root, 
+                   mesh_dump_root=opt.mesh_dump_root,
+                   transform_root=opt.transform_root,
+                   view_indices=view_indices)
+    dual_grids = dataset_utils.foreach_instance(metadata, opt.root, func, max_workers=opt.max_workers, desc='Dual griding views', timeout=300)
+    
+    # Processing summary
+    total_processed = dual_grids['_processed_count'].sum() if '_processed_count' in dual_grids.columns else 0
+    total_skipped = dual_grids['_skipped_count'].sum() if '_skipped_count' in dual_grids.columns else 0
+    print(f'\n========== Processing Summary ==========')
+    print(f'Total processed (new): {int(total_processed)}')
+    print(f'Total skipped (existing): {int(total_skipped)}')
+    print(f'Total items: {int(total_processed + total_skipped)}')
+    print(f'=========================================\n')
+    
+    if 'error' in dual_grids.columns:
+        errors = dual_grids[dual_grids['error'].notna()]
+        if len(errors) > 0:
+            with open('errors_view.txt', 'w') as f:
+                f.write('\n'.join(errors['sha256'].tolist()))
+            print(f'Errors written to errors_view.txt ({len(errors)} errors)')
+    
+    # Save metadata
+    for res in opt.resolution:
+        # Collect all view-related columns
+        view_cols = [col for col in dual_grids.columns if f'dual_grid_view' in col and f'_{res}' in col and 'converted' in col]
+        if view_cols:
+            # Save metadata for each view
+            dual_grid_metadata = dual_grids[dual_grids[view_cols].any(axis=1)]
+            if len(dual_grid_metadata) > 0:
+                # Save simplified metadata
+                cols_to_save = ['sha256'] + [col for col in dual_grids.columns if f'_{res}' in col]
+                cols_to_save = [col for col in cols_to_save if col in dual_grids.columns]
+                dual_grid_metadata[cols_to_save].to_csv(
+                    os.path.join(opt.dual_grid_root, f'dual_grid_view_{res}', 'new_records', f'part_{opt.rank}.csv'), 
+                    index=False
+                )
+    
+    print('Done!')
diff --git a/data_toolkit/dump_mesh.py b/data_toolkit/dump_mesh.py
new file mode 100644
index 0000000..93c872e
--- /dev/null
+++ b/data_toolkit/dump_mesh.py
@@ -0,0 +1,126 @@
+import os
+import shutil
+import copy
+import sys
+import importlib
+import argparse
+import pandas as pd
+from easydict import EasyDict as edict
+from functools import partial
+from subprocess import DEVNULL, call
+import numpy as np
+import tempfile
+
+
+BLENDER_LINK = 'https://ftp.halifax.rwth-aachen.de/blender/release/Blender4.5/blender-4.5.1-linux-x64.tar.xz'
+BLENDER_INSTALLATION_PATH = '/tmp'
+BLENDER_PATH = f'{BLENDER_INSTALLATION_PATH}/blender-4.5.1-linux-x64/blender'
+
+def _install_blender():
+    if not os.path.exists(BLENDER_PATH):
+        os.system('sudo apt-get update')
+        os.system('sudo apt-get install -y libxrender1 libxi6 libxkbcommon-x11-0 libsm6 libxfixes3 libgl1')
+        os.system(f'wget {BLENDER_LINK} -P {BLENDER_INSTALLATION_PATH}')
+        os.system(f'tar -xvf {BLENDER_INSTALLATION_PATH}/blender-4.5.1-linux-x64.tar.xz -C {BLENDER_INSTALLATION_PATH}')
+
+
+def _dump_mesh(file_path, sha256, root):
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        temp_path = os.path.join(tmp_dir, f'{sha256}.pickle')
+        output_path = os.path.join(root, 'mesh_dumps', f'{sha256}.pickle')
+        args = [
+            BLENDER_PATH, '-b', '-P', os.path.join(os.path.dirname(__file__), 'blender_script', 'dump_mesh.py'),
+            '--',
+            '--object', os.path.expanduser(file_path),
+            '--output_path', os.path.expanduser(temp_path)
+        ]
+        if file_path.endswith('.blend'):
+            args.insert(1, file_path)
+        
+        call(args, stdout=DEVNULL, stderr=DEVNULL)
+        
+        if os.path.exists(temp_path):
+            shutil.move(temp_path, output_path)
+            return {'sha256': sha256, 'mesh_dumped': True}
+        else:
+            if os.path.exists(temp_path + '_error.txt'):
+                with open(temp_path + '_error.txt', 'r') as f:
+                    error_msg = f.read()
+                raise ValueError(f'Failed to dump mesh. File {file_path}. Error message: {error_msg}')
+            else:
+                raise ValueError(f'Failed to dump mesh. File {file_path}.')
+
+if __name__ == '__main__':
+    dataset_utils = importlib.import_module(f'datasets.{sys.argv[1]}')
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--root', type=str, required=True,
+                        help='Directory to save the metadata')
+    parser.add_argument('--download_root', type=str, default=None,
+                        help='Directory to save the downloaded files')
+    parser.add_argument('--mesh_dump_root', type=str, default=None,
+                        help='Directory to save the mesh dumps')
+    parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
+                        help='Filter objects with aesthetic score lower than this value')
+    parser.add_argument('--instances', type=str, default=None,
+                        help='Instances to process')
+    dataset_utils.add_args(parser)
+    parser.add_argument('--rank', type=int, default=0)
+    parser.add_argument('--world_size', type=int, default=1)
+    parser.add_argument('--max_workers', type=int, default=0)
+    opt = parser.parse_args(sys.argv[2:])
+    opt = edict(vars(opt))
+    opt.download_root = opt.download_root or opt.root
+    opt.mesh_dump_root = opt.mesh_dump_root or opt.root
+
+    os.makedirs(os.path.join(opt.mesh_dump_root, 'mesh_dumps', 'new_records'), exist_ok=True)
+    
+    # install blender
+    print('Checking blender...', flush=True)
+    _install_blender()
+
+    # get file list
+    if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
+        raise ValueError('metadata.csv not found')
+    metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
+    if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256'))
+    if os.path.exists(os.path.join(opt.download_root, 'raw', 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.download_root, 'raw', 'metadata.csv')).set_index('sha256'))
+    if os.path.exists(os.path.join(opt.mesh_dump_root, 'mesh_dumps', 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.mesh_dump_root, 'mesh_dumps', 'metadata.csv')).set_index('sha256'))
+    metadata = metadata.reset_index()
+    if opt.instances is None:
+        metadata = metadata[metadata['local_path'].notna()]
+        if opt.filter_low_aesthetic_score is not None:
+            metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
+        if 'mesh_dumped' in metadata.columns:
+            metadata = metadata[metadata['mesh_dumped'] != True]
+    else:
+        if os.path.exists(opt.instances):
+            with open(opt.instances, 'r') as f:
+                instances = f.read().splitlines()
+        else:
+            instances = opt.instances.split(',')
+        metadata = metadata[metadata['sha256'].isin(instances)]
+
+    start = len(metadata) * opt.rank // opt.world_size
+    end = len(metadata) * (opt.rank + 1) // opt.world_size
+    metadata = metadata[start:end]
+    records = []
+
+    # filter out objects that are already processed
+    sha256_list = os.listdir(os.path.join(opt.mesh_dump_root, 'mesh_dumps'))
+    sha256_list = [os.path.splitext(f)[0] for f in sha256_list if f.endswith('.pickle')]
+    for sha256 in sha256_list:
+        records.append({'sha256': sha256, 'mesh_dumped': True})
+    print(f'Found {len(sha256_list)} dumped mesh')
+    metadata = metadata[~metadata['sha256'].isin(sha256_list)]
+                
+    print(f'Processing {len(metadata)} objects...')
+
+    # process objects
+    func = partial(_dump_mesh, root=opt.mesh_dump_root)
+    mesh_dumped = dataset_utils.foreach_instance(metadata, opt.download_root, func, max_workers=opt.max_workers, desc='Dumping mesh')
+    mesh_dumped = pd.concat([mesh_dumped, pd.DataFrame.from_records(records)])
+    mesh_dumped.to_csv(os.path.join(opt.mesh_dump_root, 'mesh_dumps', 'new_records', f'part_{opt.rank}.csv'), index=False)
diff --git a/data_toolkit/dump_pbr.py b/data_toolkit/dump_pbr.py
new file mode 100644
index 0000000..5198ff6
--- /dev/null
+++ b/data_toolkit/dump_pbr.py
@@ -0,0 +1,127 @@
+import os
+import shutil
+import copy
+import sys
+import importlib
+import argparse
+import pandas as pd
+from easydict import EasyDict as edict
+from functools import partial
+from subprocess import DEVNULL, call
+import numpy as np
+import tempfile
+
+
+BLENDER_LINK = 'https://ftp.halifax.rwth-aachen.de/blender/release/Blender4.5/blender-4.5.1-linux-x64.tar.xz'
+BLENDER_INSTALLATION_PATH = '/tmp'
+BLENDER_PATH = f'{BLENDER_INSTALLATION_PATH}/blender-4.5.1-linux-x64/blender'
+
+def _install_blender():
+    if not os.path.exists(BLENDER_PATH):
+        os.system('sudo apt-get update')
+        os.system('sudo apt-get install -y libxrender1 libxi6 libxkbcommon-x11-0 libsm6 libxfixes3 libgl1')
+        os.system(f'wget {BLENDER_LINK} -P {BLENDER_INSTALLATION_PATH}')
+        os.system(f'tar -xvf {BLENDER_INSTALLATION_PATH}/blender-4.5.1-linux-x64.tar.xz -C {BLENDER_INSTALLATION_PATH}')
+    os.system(f'{BLENDER_PATH} -b --python {os.path.join(os.path.dirname(__file__), "blender_script", "install_pillow.py")}')
+
+
+def _dump_pbr(file_path, sha256, root):
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        temp_path = os.path.join(tmp_dir, f'{sha256}.pickle')
+        output_path = os.path.join(root, 'pbr_dumps', f'{sha256}.pickle')
+        args = [
+            BLENDER_PATH, '-b', '-P', os.path.join(os.path.dirname(__file__), 'blender_script', 'dump_pbr.py'),
+            '--',
+            '--object', os.path.expanduser(file_path),
+            '--output_path', os.path.expanduser(temp_path)
+        ]
+        if file_path.endswith('.blend'):
+            args.insert(1, file_path)
+        
+        call(args, stdout=DEVNULL, stderr=DEVNULL)
+        
+        if os.path.exists(temp_path):
+            shutil.move(temp_path, output_path)
+            return {'sha256': sha256, 'pbr_dumped': True}
+        else:
+            if os.path.exists(temp_path + '_error.txt'):
+                with open(temp_path + '_error.txt', 'r') as f:
+                    error_msg = f.read()
+                raise ValueError(f'Failed to dump PBR. File {file_path}. Error message: {error_msg}')
+            else:
+                raise ValueError(f'Failed to dump PBR. File {file_path}.')
+
+if __name__ == '__main__':
+    dataset_utils = importlib.import_module(f'datasets.{sys.argv[1]}')
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--root', type=str, required=True,
+                        help='Directory to save the metadata')
+    parser.add_argument('--download_root', type=str, default=None,
+                        help='Directory to save the downloaded files')
+    parser.add_argument('--pbr_dump_root', type=str, default=None,
+                        help='Directory to save the mesh dumps')
+    parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
+                        help='Filter objects with aesthetic score lower than this value')
+    parser.add_argument('--instances', type=str, default=None,
+                        help='Instances to process')
+    dataset_utils.add_args(parser)
+    parser.add_argument('--rank', type=int, default=0)
+    parser.add_argument('--world_size', type=int, default=1)
+    parser.add_argument('--max_workers', type=int, default=0)
+    opt = parser.parse_args(sys.argv[2:])
+    opt = edict(vars(opt))
+    opt.download_root = opt.download_root or opt.root
+    opt.pbr_dump_root = opt.pbr_dump_root or opt.root
+
+    os.makedirs(os.path.join(opt.pbr_dump_root, 'pbr_dumps', 'new_records'), exist_ok=True)
+    
+    # install blender
+    print('Checking blender...', flush=True)
+    _install_blender()
+
+    # get file list
+    if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
+        raise ValueError('metadata.csv not found')
+    metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
+    if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256'))
+    if os.path.exists(os.path.join(opt.download_root, 'raw', 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.download_root, 'raw', 'metadata.csv')).set_index('sha256'))
+    if os.path.exists(os.path.join(opt.pbr_dump_root, 'pbr_dumps', 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.pbr_dump_root, 'pbr_dumps', 'metadata.csv')).set_index('sha256'))
+    metadata = metadata.reset_index()
+    if opt.instances is None:
+        metadata = metadata[metadata['local_path'].notna()]
+        if opt.filter_low_aesthetic_score is not None:
+            metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
+        if 'pbr_dumped' in metadata.columns:
+            metadata = metadata[metadata['pbr_dumped'] != True]
+    else:
+        if os.path.exists(opt.instances):
+            with open(opt.instances, 'r') as f:
+                instances = f.read().splitlines()
+        else:
+            instances = opt.instances.split(',')
+        metadata = metadata[metadata['sha256'].isin(instances)]
+
+    start = len(metadata) * opt.rank // opt.world_size
+    end = len(metadata) * (opt.rank + 1) // opt.world_size
+    metadata = metadata[start:end]
+    records = []
+
+    # filter out objects that are already processed
+    sha256_list = os.listdir(os.path.join(opt.pbr_dump_root, 'pbr_dumps'))
+    sha256_list = [os.path.splitext(f)[0] for f in sha256_list if f.endswith('.pickle')]
+    for sha256 in sha256_list:
+        records.append({'sha256': sha256, 'pbr_dumped': True})
+    print(f'Found {len(sha256_list)} dumped PBRs')
+    metadata = metadata[~metadata['sha256'].isin(sha256_list)]
+       
+    print(f'Processing {len(metadata)} objects...')
+
+    # process objects
+    func = partial(_dump_pbr, root=opt.pbr_dump_root)
+    pbr_dumped = dataset_utils.foreach_instance(metadata, opt.download_root, func, max_workers=opt.max_workers, desc='Dumping PBR')
+    pbr_dumped = pd.concat([pbr_dumped, pd.DataFrame.from_records(records)])
+    pbr_dumped.to_csv(os.path.join(opt.pbr_dump_root, 'pbr_dumps', 'new_records', f'part_{opt.rank}.csv'), index=False)
diff --git a/data_toolkit/encode_pbr_latent.py b/data_toolkit/encode_pbr_latent.py
new file mode 100644
index 0000000..18d5ede
--- /dev/null
+++ b/data_toolkit/encode_pbr_latent.py
@@ -0,0 +1,181 @@
+import os
+import sys
+sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
+import json
+import argparse
+import torch
+import numpy as np
+import pandas as pd
+import o_voxel
+from tqdm import tqdm
+from easydict import EasyDict as edict
+from concurrent.futures import ThreadPoolExecutor
+from queue import Queue
+
+import pixal3d.models as models
+import pixal3d.modules.sparse as sp
+
+torch.set_grad_enabled(False)
+
+def is_valid_sparse_tensor(tensor):
+    return torch.isfinite(tensor.feats).all() and torch.isfinite(tensor.coords).all()
+
+def clear_cuda_error():
+    torch.cuda.synchronize()
+    torch.cuda.empty_cache()
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--root', type=str, required=True,
+                        help='Directory to save the metadata')
+    parser.add_argument('--pbr_voxel_root', type=str, default=None,
+                        help='Directory to save the pbr voxel files')
+    parser.add_argument('--pbr_latent_root', type=str, default=None,
+                        help='Directory to save the pbr latent files')
+    parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
+                        help='Filter objects with aesthetic score lower than this value')
+    parser.add_argument('--resolution', type=int, default=1024,
+                        help='Sparse voxel resolution')
+    parser.add_argument('--enc_pretrained', type=str, default='microsoft/TRELLIS.2-4B/ckpts/tex_enc_next_dc_f16c32_fp16',
+                        help='Pretrained encoder model')
+    parser.add_argument('--model_root', type=str,
+                        help='Root directory of models')
+    parser.add_argument('--enc_model', type=str,
+                        help='Encoder model. if specified, use this model instead of pretrained model')
+    parser.add_argument('--ckpt', type=str,
+                        help='Checkpoint to load')
+    parser.add_argument('--instances', type=str, default=None,
+                        help='Instances to process')
+    parser.add_argument('--rank', type=int, default=0)
+    parser.add_argument('--world_size', type=int, default=1)
+    opt = parser.parse_args()
+    opt = edict(vars(opt))
+    opt.pbr_voxel_root = opt.pbr_voxel_root or opt.root
+    opt.pbr_latent_root = opt.pbr_latent_root or opt.root
+
+    if opt.enc_model is None:
+        latent_name = f'{opt.enc_pretrained.split("/")[-1]}_{opt.resolution}'
+        encoder = models.from_pretrained(opt.enc_pretrained).eval().cuda()
+    else:
+        latent_name = f'{opt.enc_model.split("/")[-1]}_{opt.ckpt}_{opt.resolution}'
+        cfg = edict(json.load(open(os.path.join(opt.model_root, opt.enc_model, 'config.json'), 'r')))
+        encoder = getattr(models, cfg.models.encoder.name)(**cfg.models.encoder.args).cuda()
+        ckpt_path = os.path.join(opt.model_root, opt.enc_model, 'ckpts', f'encoder_{opt.ckpt}.pt')
+        encoder.load_state_dict(torch.load(ckpt_path), strict=False)
+        encoder.eval()
+        print(f'Loaded model from {ckpt_path}')
+    
+    os.makedirs(os.path.join(opt.pbr_latent_root, 'pbr_latents', latent_name, 'new_records'), exist_ok=True)
+    
+    # get file list
+    if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
+        raise ValueError('metadata.csv not found')
+    metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
+    if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256'))
+    if os.path.exists(os.path.join(opt.pbr_voxel_root, f'pbr_voxels_{opt.resolution}', 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.pbr_voxel_root, f'pbr_voxels_{opt.resolution}','metadata.csv')).set_index('sha256'))
+    if os.path.exists(os.path.join(opt.pbr_latent_root, 'pbr_latents', latent_name, 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.pbr_latent_root, 'pbr_latents', latent_name,'metadata.csv')).set_index('sha256'))
+    metadata = metadata.reset_index()
+    if opt.instances is None:
+        if opt.filter_low_aesthetic_score is not None:
+            metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
+        metadata = metadata[metadata['pbr_voxelized'] == True]
+        if 'pbr_latent_encoded' in metadata.columns:
+            metadata = metadata[metadata['pbr_latent_encoded'] != True]
+    else:
+        if os.path.exists(opt.instances):
+            with open(opt.instances, 'r') as f:
+                instances = f.read().splitlines()
+        else:
+            instances = opt.instances.split(',')
+        metadata = metadata[metadata['sha256'].isin(instances)]
+
+    start = len(metadata) * opt.rank // opt.world_size
+    end = len(metadata) * (opt.rank + 1) // opt.world_size
+    metadata = metadata[start:end]
+    records = []
+    
+    # filter out objects that are already processed
+    with ThreadPoolExecutor(max_workers=os.cpu_count()) as executor, \
+        tqdm(total=len(metadata), desc="Filtering existing objects") as pbar:
+        def check_sha256(sha256):
+            if os.path.exists(os.path.join(opt.pbr_latent_root, 'pbr_latents', latent_name, f'{sha256}.npz')):
+                coords = np.load(os.path.join(opt.pbr_latent_root, 'pbr_latents', latent_name, f'{sha256}.npz'))['coords']
+                records.append({'sha256': sha256, 'pbr_latent_encoded': True, 'pbr_latent_tokens': coords.shape[0]})
+            pbar.update()
+        executor.map(check_sha256, metadata['sha256'].values)
+        executor.shutdown(wait=True)
+    existing_sha256 = set(r['sha256'] for r in records)
+    print(f'Found {len(existing_sha256)} processed objects')
+    metadata = metadata[~metadata['sha256'].isin(existing_sha256)]
+        
+    print(f'Processing {len(metadata)} objects...')
+    
+    sha256s = list(metadata['sha256'].values)
+    load_queue = Queue(maxsize=32)
+    with ThreadPoolExecutor(max_workers=32) as loader_executor, \
+         ThreadPoolExecutor(max_workers=32) as saver_executor:
+
+        def loader(sha256):
+            try:
+                attrs = ['base_color', 'metallic', 'roughness', 'alpha']                
+                coords, attr = o_voxel.io.read_vxz(
+                    os.path.join(opt.pbr_voxel_root, f'pbr_voxels_{opt.resolution}', f'{sha256}.vxz'),
+                    num_threads=4
+                )
+                feats = torch.concat([attr[k] for k in attrs], dim=-1) / 255.0 * 2 - 1
+                x = sp.SparseTensor(
+                    feats.float(),
+                    torch.cat([torch.zeros_like(coords[:, 0:1]), coords], dim=-1),
+                )
+                load_queue.put((sha256, x))
+            except Exception as e:
+                print(f"[Loader Error] {sha256}: {e}")
+                load_queue.put((sha256, None))
+
+        loader_executor.map(loader, sha256s)
+        
+        def saver(sha256, pack):
+            save_path = os.path.join(opt.pbr_latent_root, 'pbr_latents', latent_name, f'{sha256}.npz')
+            np.savez_compressed(save_path, **pack)
+            records.append({'sha256': sha256, 'pbr_latent_encoded': True, 'pbr_latent_tokens': pack['coords'].shape[0]})
+            
+        for _ in tqdm(range(len(sha256s)), desc="Extracting latents"):
+            try:
+                sha256, voxels = load_queue.get()
+                if voxels is None:
+                    print(f"[Skip] {sha256}: Failed to load input")
+                    continue
+                
+                num_voxels = voxels.feats.shape[0]
+
+                # NaN/Inf
+                if not (is_valid_sparse_tensor(voxels)):
+                    print(f"[Skip] {sha256}: NaN/Inf in input")
+                    continue
+
+                z = encoder(voxels.cuda())
+                torch.cuda.synchronize()
+
+                if not torch.isfinite(z.feats).all():
+                    print(f"[Skip] {sha256}: Non-finite latent in z.feats")
+                    clear_cuda_error()
+                    continue
+
+                pack = {
+                    'feats': z.feats.cpu().numpy().astype(np.float32),
+                    'coords': z.coords[:, 1:].cpu().numpy().astype(np.uint8),
+                }
+                saver_executor.submit(saver, sha256, pack)
+
+            except Exception as e:
+                print(f"[Error] {sha256} ({num_voxels} voxels): {e}")
+                clear_cuda_error()
+                continue
+            
+        saver_executor.shutdown(wait=True)
+        
+    records = pd.DataFrame.from_records(records)
+    records.to_csv(os.path.join(opt.pbr_latent_root, 'pbr_latents', latent_name, 'new_records', f'part_{opt.rank}.csv'), index=False)
diff --git a/data_toolkit/encode_pbr_latent_view.py b/data_toolkit/encode_pbr_latent_view.py
new file mode 100644
index 0000000..75928cf
--- /dev/null
+++ b/data_toolkit/encode_pbr_latent_view.py
@@ -0,0 +1,271 @@
+import os
+import sys
+sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
+import json
+import argparse
+import shutil
+import torch
+import numpy as np
+import pandas as pd
+import o_voxel
+from tqdm import tqdm
+from easydict import EasyDict as edict
+from concurrent.futures import ThreadPoolExecutor
+from queue import Queue
+from utils import parse_view_indices
+
+import pixal3d.models as models
+import pixal3d.modules.sparse as sp
+
+torch.set_grad_enabled(False)
+
+
+def is_valid_sparse_tensor(tensor):
+    return torch.isfinite(tensor.feats).all() and torch.isfinite(tensor.coords).all()
+
+def clear_cuda_error():
+    torch.cuda.synchronize()
+    torch.cuda.empty_cache()
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--root', type=str, required=True,
+                        help='Directory to save the metadata')
+    parser.add_argument('--pbr_voxel_root', type=str, default=None,
+                        help='Directory containing the pbr voxels')
+    parser.add_argument('--pbr_latent_root', type=str, default=None,
+                        help='Directory to save the pbr latent files')
+    parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
+                        help='Filter objects with aesthetic score lower than this value')
+    parser.add_argument('--resolution', type=int, default=1024,
+                        help='Sparse voxel resolution')
+    parser.add_argument('--enc_pretrained', type=str, default='microsoft/TRELLIS.2-4B/ckpts/tex_enc_next_dc_f16c32_fp16',
+                        help='Pretrained encoder model')
+    parser.add_argument('--model_root', type=str,
+                        help='Root directory of models')
+    parser.add_argument('--enc_model', type=str,
+                        help='Encoder model. if specified, use this model instead of pretrained model')
+    parser.add_argument('--ckpt', type=str,
+                        help='Checkpoint to load')
+    parser.add_argument('--instances', type=str, default=None,
+                        help='Instances to process')
+    parser.add_argument('--view_indices', type=str, default=None,
+                        help='View indices to process, e.g., "0,1,2" or "0-5". None for all views')
+    parser.add_argument('--num_views', type=int, default=24,
+                        help='Total number of views (used when view_indices is None)')
+    parser.add_argument('--rank', type=int, default=0)
+    parser.add_argument('--world_size', type=int, default=1)
+    opt = parser.parse_args()
+    opt = edict(vars(opt))
+    opt.pbr_voxel_root = opt.pbr_voxel_root or opt.root
+    opt.pbr_latent_root = opt.pbr_latent_root or opt.root
+
+    # Parse view_indices
+    view_indices = parse_view_indices(opt.view_indices)
+    if view_indices is None:
+        view_indices = list(range(opt.num_views))
+    
+    print(f'View indices to process: {view_indices}')
+
+    if opt.enc_model is None:
+        latent_name = f'{opt.enc_pretrained.split("/")[-1]}_{opt.resolution}'
+        encoder = models.from_pretrained(opt.enc_pretrained).eval().cuda()
+    else:
+        latent_name = f'{opt.enc_model.split("/")[-1]}_{opt.ckpt}_{opt.resolution}'
+        cfg = edict(json.load(open(os.path.join(opt.model_root, opt.enc_model, 'config.json'), 'r')))
+        encoder = getattr(models, cfg.models.encoder.name)(**cfg.models.encoder.args).cuda()
+        ckpt_path = os.path.join(opt.model_root, opt.enc_model, 'ckpts', f'encoder_{opt.ckpt}.pt')
+        encoder.load_state_dict(torch.load(ckpt_path), strict=False)
+        encoder.eval()
+        print(f'Loaded model from {ckpt_path}')
+    
+    # Multi-view latent output directory
+    latent_view_name = f'{latent_name}_view_fix'
+    os.makedirs(os.path.join(opt.pbr_latent_root, 'pbr_latents', latent_view_name, 'new_records'), exist_ok=True)
+    
+    # Get file list
+    if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
+        raise ValueError('metadata.csv not found')
+    metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
+    if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')):
+        aesthetic_metadata = pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256')
+        metadata = metadata.join(aesthetic_metadata, how='left', rsuffix='_aesthetic')
+    
+    # Check pbr_voxels_view_fix metadata
+    pbr_voxel_view_path = os.path.join(opt.pbr_voxel_root, f'pbr_voxels_view_fix_{opt.resolution}', 'metadata.csv')
+    if os.path.exists(pbr_voxel_view_path):
+        pbr_voxel_metadata = pd.read_csv(pbr_voxel_view_path).set_index('sha256')
+        metadata = metadata.join(pbr_voxel_metadata, how='left', rsuffix='_pbr_voxel')
+    
+    # Check pbr_latent_view metadata (used to skip already completed tasks)
+    pbr_latent_view_metadata_path = os.path.join(opt.pbr_latent_root, 'pbr_latents', latent_view_name, 'metadata.csv')
+    if os.path.exists(pbr_latent_view_metadata_path):
+        pbr_latent_view_metadata = pd.read_csv(pbr_latent_view_metadata_path).set_index('sha256')
+        metadata = metadata.join(pbr_latent_view_metadata, how='left', rsuffix='_pbr_latent_view')
+        print(f'Loaded pbr_latent_view metadata with {len(pbr_latent_view_metadata)} records')
+    else:
+        print(f'Warning: pbr_latent_view metadata not found at {pbr_latent_view_metadata_path}')
+    
+    metadata = metadata.reset_index()
+    
+    if opt.instances is None:
+        if opt.filter_low_aesthetic_score is not None:
+            metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
+        
+        # Filter to objects that have pbr_voxels_view_fix data
+        # Use first view as indicator
+        first_view_col = f'pbr_voxelized_view_fix{view_indices[0]:02d}_{opt.resolution}'
+        if first_view_col in metadata.columns:
+            metadata = metadata[metadata[first_view_col] == True]
+        else:
+            print(f'Warning: Column {first_view_col} not found in metadata, will check files directly')
+    else:
+        if os.path.exists(opt.instances):
+            with open(opt.instances, 'r') as f:
+                instances = f.read().splitlines()
+        else:
+            instances = opt.instances.split(',')
+        metadata = metadata[metadata['sha256'].isin(instances)]
+
+    records = []
+    
+    # Build task list: (sha256, view_idx), filter already completed tasks via metadata
+    all_tasks = []
+    skipped_count = 0
+    
+    # Pre-fetch completion status columns for each view
+    encoded_cols = {view_idx: f'pbr_latent_view{view_idx:02d}_encoded' for view_idx in view_indices}
+    
+    for _, row in metadata.iterrows():
+        sha256 = row['sha256']
+        for view_idx in view_indices:
+            encoded_col = encoded_cols[view_idx]
+            # Check if already marked as completed in metadata
+            if encoded_col in metadata.columns and row.get(encoded_col, False) == True:
+                skipped_count += 1
+                continue
+            all_tasks.append((sha256, view_idx))
+    
+    # Split tasks by rank after filtering completed ones
+    start = len(all_tasks) * opt.rank // opt.world_size
+    end = len(all_tasks) * (opt.rank + 1) // opt.world_size
+    tasks = all_tasks[start:end]
+    
+    print(f'Total tasks: {len(all_tasks) + skipped_count}, Already done (from metadata): {skipped_count}, To process (all ranks): {len(all_tasks)}, This rank: {len(tasks)}')
+    
+    load_queue = Queue(maxsize=32)
+    
+    with ThreadPoolExecutor(max_workers=32) as loader_executor, \
+         ThreadPoolExecutor(max_workers=32) as saver_executor:
+
+        def loader(task):
+            sha256, view_idx = task
+            try:
+                # Check if output file already exists, skip if so (but still record)
+                output_path = os.path.join(
+                    opt.pbr_latent_root, 
+                    'pbr_latents', 
+                    latent_view_name, 
+                    sha256, 
+                    f'view{view_idx:02d}.npz'
+                )
+                if os.path.exists(output_path):
+                    try:
+                        data = np.load(output_path)
+                        num_tokens = data['coords'].shape[0]
+                    except Exception:
+                        num_tokens = -1
+                    records.append({
+                        'sha256': sha256,
+                        f'pbr_latent_view{view_idx:02d}_encoded': True,
+                        f'pbr_latent_view{view_idx:02d}_tokens': num_tokens,
+                    })
+                    load_queue.put((sha256, view_idx, None))
+                    return
+                
+                # pbr_voxels_view_fix path: pbr_voxels_view_fix_{res}/{sha256}/view{idx:02d}.vxz
+                vxz_path = os.path.join(
+                    opt.pbr_voxel_root, 
+                    f'pbr_voxels_view_fix_{opt.resolution}', 
+                    sha256, 
+                    f'view{view_idx:02d}.vxz'
+                )
+                
+                if not os.path.exists(vxz_path):
+                    print(f"[Loader Skip] {sha256}/view{view_idx:02d}: vxz file not found")
+                    load_queue.put((sha256, view_idx, None))
+                    return
+                
+                attrs = ['base_color', 'metallic', 'roughness', 'alpha']
+                coords, attr = o_voxel.io.read_vxz(vxz_path, num_threads=4)
+                feats = torch.concat([attr[k] for k in attrs], dim=-1) / 255.0 * 2 - 1
+                x = sp.SparseTensor(
+                    feats.float(),
+                    torch.cat([torch.zeros_like(coords[:, 0:1]), coords], dim=-1),
+                )
+                load_queue.put((sha256, view_idx, x))
+            except Exception as e:
+                print(f"[Loader Error] {sha256}/view{view_idx:02d}: {e}")
+                load_queue.put((sha256, view_idx, None))
+
+        loader_executor.map(loader, tasks)
+        
+        def saver(sha256, view_idx, pack):
+            sha256_dir = os.path.join(opt.pbr_latent_root, 'pbr_latents', latent_view_name, sha256)
+            os.makedirs(sha256_dir, exist_ok=True)
+            save_path = os.path.join(sha256_dir, f'view{view_idx:02d}.npz')
+            np.savez_compressed(save_path, **pack)
+            
+            # Copy scale json from pbr_voxels_view_fix
+            src_scale_path = os.path.join(
+                opt.pbr_voxel_root,
+                f'pbr_voxels_view_fix_{opt.resolution}',
+                sha256,
+                f'view{view_idx:02d}_scale.json'
+            )
+            dst_scale_path = os.path.join(sha256_dir, f'view{view_idx:02d}_scale.json')
+            if os.path.exists(src_scale_path):
+                shutil.copy2(src_scale_path, dst_scale_path)
+            
+            records.append({
+                'sha256': sha256,
+                f'pbr_latent_view{view_idx:02d}_encoded': True,
+                f'pbr_latent_view{view_idx:02d}_tokens': pack['coords'].shape[0]
+            })
+            
+        for _ in tqdm(range(len(tasks)), desc=f"Extracting {os.path.basename(opt.root)} PBR view latents (res={opt.resolution})"):
+            try:
+                sha256, view_idx, voxels = load_queue.get()
+                if voxels is None:
+                    continue
+                
+                num_voxels = voxels.feats.shape[0]
+
+                # NaN/Inf check
+                if not is_valid_sparse_tensor(voxels):
+                    print(f"[Skip] {sha256}/view{view_idx:02d}: NaN/Inf in input")
+                    continue
+
+                z = encoder(voxels.cuda())
+                torch.cuda.synchronize()
+
+                if not torch.isfinite(z.feats).all():
+                    print(f"[Skip] {sha256}/view{view_idx:02d}: Non-finite latent in z.feats")
+                    clear_cuda_error()
+                    continue
+
+                pack = {
+                    'feats': z.feats.cpu().numpy().astype(np.float32),
+                    'coords': z.coords[:, 1:].cpu().numpy().astype(np.uint8),
+                }
+                saver_executor.submit(saver, sha256, view_idx, pack)
+
+            except Exception as e:
+                print(f"[Error] {sha256}/view{view_idx:02d} ({num_voxels} voxels): {e}")
+                clear_cuda_error()
+                continue
+            
+        saver_executor.shutdown(wait=True)
+        
+    records = pd.DataFrame.from_records(records)
+    records.to_csv(os.path.join(opt.pbr_latent_root, 'pbr_latents', latent_view_name, 'new_records', f'part_{opt.rank}.csv'), index=False)
diff --git a/data_toolkit/encode_shape_latent.py b/data_toolkit/encode_shape_latent.py
new file mode 100644
index 0000000..07150b0
--- /dev/null
+++ b/data_toolkit/encode_shape_latent.py
@@ -0,0 +1,184 @@
+import os
+import sys
+sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
+import json
+import argparse
+import torch
+import numpy as np
+import pandas as pd
+import o_voxel
+from tqdm import tqdm
+from easydict import EasyDict as edict
+from concurrent.futures import ThreadPoolExecutor
+from queue import Queue
+
+import pixal3d.models as models
+import pixal3d.modules.sparse as sp
+
+torch.set_grad_enabled(False)
+
+def is_valid_sparse_tensor(tensor):
+    return torch.isfinite(tensor.feats).all() and torch.isfinite(tensor.coords).all()
+
+def clear_cuda_error():
+    torch.cuda.synchronize()
+    torch.cuda.empty_cache()
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--root', type=str, required=True,
+                        help='Directory to save the metadata')
+    parser.add_argument('--dual_grid_root', type=str, default=None,
+                        help='Directory to save the dual grids')
+    parser.add_argument('--shape_latent_root', type=str, default=None,
+                        help='Directory to save the shape latent files')
+    parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
+                        help='Filter objects with aesthetic score lower than this value')
+    parser.add_argument('--resolution', type=int, default=1024,
+                        help='Sparse voxel resolution')
+    parser.add_argument('--enc_pretrained', type=str, default='microsoft/TRELLIS.2-4B/ckpts/shape_enc_next_dc_f16c32_fp16',
+                        help='Pretrained encoder model')
+    parser.add_argument('--model_root', type=str,
+                        help='Root directory of models')
+    parser.add_argument('--enc_model', type=str,
+                        help='Encoder model. if specified, use this model instead of pretrained model')
+    parser.add_argument('--ckpt', type=str,
+                        help='Checkpoint to load')
+    parser.add_argument('--instances', type=str, default=None,
+                        help='Instances to process')
+    parser.add_argument('--rank', type=int, default=0)
+    parser.add_argument('--world_size', type=int, default=1)
+    opt = parser.parse_args()
+    opt = edict(vars(opt))
+    opt.dual_grid_root = opt.dual_grid_root or opt.root
+    opt.shape_latent_root = opt.shape_latent_root or opt.root
+
+    if opt.enc_model is None:
+        latent_name = f'{opt.enc_pretrained.split("/")[-1]}_{opt.resolution}'
+        encoder = models.from_pretrained(opt.enc_pretrained).eval().cuda()
+    else:
+        latent_name = f'{opt.enc_model.split("/")[-1]}_{opt.ckpt}_{opt.resolution}'
+        cfg = edict(json.load(open(os.path.join(opt.model_root, opt.enc_model, 'config.json'), 'r')))
+        encoder = getattr(models, cfg.models.encoder.name)(**cfg.models.encoder.args).cuda()
+        ckpt_path = os.path.join(opt.model_root, opt.enc_model, 'ckpts', f'encoder_{opt.ckpt}.pt')
+        encoder.load_state_dict(torch.load(ckpt_path), strict=False)
+        encoder.eval()
+        print(f'Loaded model from {ckpt_path}')
+    
+    os.makedirs(os.path.join(opt.shape_latent_root, 'shape_latents', latent_name, 'new_records'), exist_ok=True)
+    
+    # get file list
+    if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
+        raise ValueError('metadata.csv not found')
+    metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
+    if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256'))
+    if os.path.exists(os.path.join(opt.dual_grid_root, f'dual_grid_{opt.resolution}', 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.dual_grid_root, f'dual_grid_{opt.resolution}','metadata.csv')).set_index('sha256'))
+    if os.path.exists(os.path.join(opt.shape_latent_root, 'shape_latents', latent_name, 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.shape_latent_root, 'shape_latents', latent_name,'metadata.csv')).set_index('sha256'))
+    metadata = metadata.reset_index()
+    if opt.instances is None:
+        if opt.filter_low_aesthetic_score is not None:
+            metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
+        metadata = metadata[metadata['dual_grid_converted'] == True]
+        if 'shape_latent_encoded' in metadata.columns:
+            metadata = metadata[metadata['shape_latent_encoded'] != True]
+    else:
+        if os.path.exists(opt.instances):
+            with open(opt.instances, 'r') as f:
+                instances = f.read().splitlines()
+        else:
+            instances = opt.instances.split(',')
+        metadata = metadata[metadata['sha256'].isin(instances)]
+
+    start = len(metadata) * opt.rank // opt.world_size
+    end = len(metadata) * (opt.rank + 1) // opt.world_size
+    metadata = metadata[start:end]
+    records = []
+    
+    # filter out objects that are already processed
+    with ThreadPoolExecutor(max_workers=os.cpu_count()) as executor, \
+        tqdm(total=len(metadata), desc="Filtering existing objects") as pbar:
+        def check_sha256(sha256):
+            if os.path.exists(os.path.join(opt.shape_latent_root, 'shape_latents', latent_name, f'{sha256}.npz')):
+                coords = np.load(os.path.join(opt.shape_latent_root, 'shape_latents', latent_name, f'{sha256}.npz'))['coords']
+                records.append({'sha256': sha256, 'shape_latent_encoded': True, 'shape_latent_tokens': coords.shape[0]})
+            pbar.update()
+        executor.map(check_sha256, metadata['sha256'].values)
+        executor.shutdown(wait=True)
+    existing_sha256 = set(r['sha256'] for r in records)
+    print(f'Found {len(existing_sha256)} processed objects')
+    metadata = metadata[~metadata['sha256'].isin(existing_sha256)]
+        
+    print(f'Processing {len(metadata)} objects...')
+    
+    sha256s = list(metadata['sha256'].values)
+    load_queue = Queue(maxsize=32)
+    with ThreadPoolExecutor(max_workers=32) as loader_executor, \
+         ThreadPoolExecutor(max_workers=32) as saver_executor:
+
+        def loader(sha256):
+            try:
+                coords, attr = o_voxel.io.read_vxz(
+                    os.path.join(opt.dual_grid_root, f'dual_grid_{opt.resolution}', f'{sha256}.vxz'),
+                    num_threads=4
+                )
+                vertices = sp.SparseTensor(
+                    (attr['vertices'] / 255.0).float(),
+                    torch.cat([torch.zeros_like(coords[:, 0:1]), coords], dim=-1),
+                )
+                intersected = vertices.replace(torch.cat([
+                    attr['intersected'] % 2,
+                    attr['intersected'] // 2 % 2,
+                    attr['intersected'] // 4 % 2,
+                ], dim=-1).bool())
+                load_queue.put((sha256, vertices, intersected))
+            except Exception as e:
+                print(f"[Loader Error] {sha256}: {e}")
+                load_queue.put((sha256, None, None))
+
+        loader_executor.map(loader, sha256s)
+        
+        def saver(sha256, pack):
+            save_path = os.path.join(opt.shape_latent_root, 'shape_latents', latent_name, f'{sha256}.npz')
+            np.savez_compressed(save_path, **pack)
+            records.append({'sha256': sha256, 'shape_latent_encoded': True, 'shape_latent_tokens': pack['coords'].shape[0]})
+            
+        for _ in tqdm(range(len(sha256s)), desc="Extracting latents"):
+            try:
+                sha256, vertices, intersected = load_queue.get()
+                if vertices is None or intersected is None:
+                    print(f"[Skip] {sha256}: Failed to load input")
+                    continue
+                
+                num_voxels = vertices.feats.shape[0]
+
+                # NaN/Inf
+                if not (is_valid_sparse_tensor(vertices) and is_valid_sparse_tensor(intersected)):
+                    print(f"[Skip] {sha256}: NaN/Inf in input")
+                    continue
+
+                z = encoder(vertices.cuda(), intersected.cuda())
+                torch.cuda.synchronize()
+
+                if not torch.isfinite(z.feats).all():
+                    print(f"[Skip] {sha256}: Non-finite latent in z.feats")
+                    clear_cuda_error()
+                    continue
+
+                pack = {
+                    'feats': z.feats.cpu().numpy().astype(np.float32),
+                    'coords': z.coords[:, 1:].cpu().numpy().astype(np.uint8),
+                }
+                saver_executor.submit(saver, sha256, pack)
+
+            except Exception as e:
+                print(f"[Error] {sha256} ({num_voxels} voxels): {e}")
+                clear_cuda_error()
+                continue
+            
+        saver_executor.shutdown(wait=True)
+        
+    records = pd.DataFrame.from_records(records)
+    records.to_csv(os.path.join(opt.shape_latent_root, 'shape_latents', latent_name, 'new_records', f'part_{opt.rank}.csv'), index=False)
diff --git a/data_toolkit/encode_shape_latent_view.py b/data_toolkit/encode_shape_latent_view.py
new file mode 100644
index 0000000..1d36009
--- /dev/null
+++ b/data_toolkit/encode_shape_latent_view.py
@@ -0,0 +1,262 @@
+import os
+import sys
+sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
+import json
+import argparse
+import shutil
+import torch
+import numpy as np
+import pandas as pd
+import o_voxel
+from tqdm import tqdm
+from easydict import EasyDict as edict
+from concurrent.futures import ThreadPoolExecutor
+from queue import Queue
+from utils import parse_view_indices
+
+import pixal3d.models as models
+import pixal3d.modules.sparse as sp
+
+torch.set_grad_enabled(False)
+
+
+def is_valid_sparse_tensor(tensor):
+    return torch.isfinite(tensor.feats).all() and torch.isfinite(tensor.coords).all()
+
+def clear_cuda_error():
+    torch.cuda.synchronize()
+    torch.cuda.empty_cache()
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--root', type=str, required=True,
+                        help='Directory to save the metadata')
+    parser.add_argument('--dual_grid_root', type=str, default=None,
+                        help='Directory containing the dual grids')
+    parser.add_argument('--shape_latent_root', type=str, default=None,
+                        help='Directory to save the shape latent files')
+    parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
+                        help='Filter objects with aesthetic score lower than this value')
+    parser.add_argument('--resolution', type=int, default=1024,
+                        help='Sparse voxel resolution')
+    parser.add_argument('--enc_pretrained', type=str, default='microsoft/TRELLIS.2-4B/ckpts/shape_enc_next_dc_f16c32_fp16',
+                        help='Pretrained encoder model')
+    parser.add_argument('--model_root', type=str,
+                        help='Root directory of models')
+    parser.add_argument('--enc_model', type=str,
+                        help='Encoder model. if specified, use this model instead of pretrained model')
+    parser.add_argument('--ckpt', type=str,
+                        help='Checkpoint to load')
+    parser.add_argument('--instances', type=str, default=None,
+                        help='Instances to process')
+    parser.add_argument('--view_indices', type=str, default=None,
+                        help='View indices to process, e.g., "0,1,2" or "0-5". None for all views')
+    parser.add_argument('--num_views', type=int, default=24,
+                        help='Total number of views (used when view_indices is None)')
+    parser.add_argument('--rank', type=int, default=0)
+    parser.add_argument('--world_size', type=int, default=1)
+    opt = parser.parse_args()
+    opt = edict(vars(opt))
+    opt.dual_grid_root = opt.dual_grid_root or opt.root
+    opt.shape_latent_root = opt.shape_latent_root or opt.root
+
+    # Parse view_indices
+    view_indices = parse_view_indices(opt.view_indices)
+    if view_indices is None:
+        view_indices = list(range(opt.num_views))
+    
+    print(f'View indices to process: {view_indices}')
+
+    if opt.enc_model is None:
+        latent_name = f'{opt.enc_pretrained.split("/")[-1]}_{opt.resolution}'
+        encoder = models.from_pretrained(opt.enc_pretrained).eval().cuda()
+    else:
+        latent_name = f'{opt.enc_model.split("/")[-1]}_{opt.ckpt}_{opt.resolution}'
+        cfg = edict(json.load(open(os.path.join(opt.model_root, opt.enc_model, 'config.json'), 'r')))
+        encoder = getattr(models, cfg.models.encoder.name)(**cfg.models.encoder.args).cuda()
+        ckpt_path = os.path.join(opt.model_root, opt.enc_model, 'ckpts', f'encoder_{opt.ckpt}.pt')
+        encoder.load_state_dict(torch.load(ckpt_path), strict=False)
+        encoder.eval()
+        print(f'Loaded model from {ckpt_path}')
+    
+    # Multi-view latent output directory
+    latent_view_name = f'{latent_name}_view'
+    os.makedirs(os.path.join(opt.shape_latent_root, 'shape_latents', latent_view_name, 'new_records'), exist_ok=True)
+    
+    # Get file list
+    if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
+        raise ValueError('metadata.csv not found')
+    metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
+    if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')):
+        aesthetic_metadata = pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256')
+        metadata = metadata.join(aesthetic_metadata, how='left', rsuffix='_aesthetic')
+    
+    # Check dual_grid_view metadata
+    dual_grid_view_path = os.path.join(opt.dual_grid_root, f'dual_grid_view_{opt.resolution}', 'metadata.csv')
+    if os.path.exists(dual_grid_view_path):
+        dual_grid_metadata = pd.read_csv(dual_grid_view_path).set_index('sha256')
+        metadata = metadata.join(dual_grid_metadata, how='left', rsuffix='_dual_grid')
+    
+    # Check shape_latent_view metadata (used to skip already completed tasks)
+    shape_latent_view_metadata_path = os.path.join(opt.shape_latent_root, 'shape_latents', latent_view_name, 'metadata.csv')
+    if os.path.exists(shape_latent_view_metadata_path):
+        shape_latent_view_metadata = pd.read_csv(shape_latent_view_metadata_path).set_index('sha256')
+        metadata = metadata.join(shape_latent_view_metadata, how='left', rsuffix='_shape_latent_view')
+        print(f'Loaded shape_latent_view metadata with {len(shape_latent_view_metadata)} records')
+    else:
+        print(f'Warning: shape_latent_view metadata not found at {shape_latent_view_metadata_path}')
+    
+    metadata = metadata.reset_index()
+    
+    if opt.instances is None:
+        if opt.filter_low_aesthetic_score is not None:
+            metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
+        
+        # Filter to objects that have dual_grid_view data
+        # Use first view as indicator
+        first_view_col = f'dual_grid_view{view_indices[0]:02d}_converted_{opt.resolution}'
+        if first_view_col in metadata.columns:
+            metadata = metadata[metadata[first_view_col] == True]
+        else:
+            print(f'Warning: Column {first_view_col} not found in metadata, will check files directly')
+    else:
+        if os.path.exists(opt.instances):
+            with open(opt.instances, 'r') as f:
+                instances = f.read().splitlines()
+        else:
+            instances = opt.instances.split(',')
+        metadata = metadata[metadata['sha256'].isin(instances)]
+
+    start = len(metadata) * opt.rank // opt.world_size
+    end = len(metadata) * (opt.rank + 1) // opt.world_size
+    metadata = metadata[start:end]
+    records = []
+    
+    # Build task list: (sha256, view_idx), filter already completed tasks via metadata
+    tasks = []
+    skipped_count = 0
+    
+    # Pre-fetch completion status columns for each view
+    encoded_cols = {view_idx: f'shape_latent_view{view_idx:02d}_encoded' for view_idx in view_indices}
+    
+    for _, row in metadata.iterrows():
+        sha256 = row['sha256']
+        for view_idx in view_indices:
+            encoded_col = encoded_cols[view_idx]
+            # Check if already marked as completed in metadata
+            if encoded_col in metadata.columns and row.get(encoded_col, False) == True:
+                skipped_count += 1
+                continue
+            tasks.append((sha256, view_idx))
+    
+    print(f'Total tasks: {len(tasks) + skipped_count}, Already done (from metadata): {skipped_count}, To process: {len(tasks)}')
+    
+    load_queue = Queue(maxsize=32)
+    
+    with ThreadPoolExecutor(max_workers=32) as loader_executor, \
+         ThreadPoolExecutor(max_workers=32) as saver_executor:
+
+        def loader(task):
+            sha256, view_idx = task
+            try:
+                # Check if output file already exists, skip if so
+                output_path = os.path.join(
+                    opt.shape_latent_root, 
+                    'shape_latents', 
+                    latent_view_name, 
+                    sha256, 
+                    f'view{view_idx:02d}.npz'
+                )
+                if os.path.exists(output_path):
+                    load_queue.put((sha256, view_idx, None, None))
+                    return
+                
+                # dual_grid_view path: dual_grid_view_{res}/{sha256}/view{idx:02d}.vxz
+                vxz_path = os.path.join(
+                    opt.dual_grid_root, 
+                    f'dual_grid_view_{opt.resolution}', 
+                    sha256, 
+                    f'view{view_idx:02d}.vxz'
+                )
+                
+                if not os.path.exists(vxz_path):
+                    print(f"[Loader Skip] {sha256}/view{view_idx:02d}: vxz file not found")
+                    load_queue.put((sha256, view_idx, None, None))
+                    return
+                
+                coords, attr = o_voxel.io.read_vxz(vxz_path, num_threads=4)
+                vertices = sp.SparseTensor(
+                    (attr['vertices'] / 255.0).float(),
+                    torch.cat([torch.zeros_like(coords[:, 0:1]), coords], dim=-1),
+                )
+                intersected = vertices.replace(torch.cat([
+                    attr['intersected'] % 2,
+                    attr['intersected'] // 2 % 2,
+                    attr['intersected'] // 4 % 2,
+                ], dim=-1).bool())
+                load_queue.put((sha256, view_idx, vertices, intersected))
+            except Exception as e:
+                print(f"[Loader Error] {sha256}/view{view_idx:02d}: {e}")
+                load_queue.put((sha256, view_idx, None, None))
+
+        loader_executor.map(loader, tasks)
+        
+        def saver(sha256, view_idx, pack):
+            sha256_dir = os.path.join(opt.shape_latent_root, 'shape_latents', latent_view_name, sha256)
+            os.makedirs(sha256_dir, exist_ok=True)
+            save_path = os.path.join(sha256_dir, f'view{view_idx:02d}.npz')
+            np.savez_compressed(save_path, **pack)
+            
+            # Copy scale json from dual_grid_view
+            src_scale_path = os.path.join(
+                opt.dual_grid_root,
+                f'dual_grid_view_{opt.resolution}',
+                sha256,
+                f'view{view_idx:02d}_scale.json'
+            )
+            dst_scale_path = os.path.join(sha256_dir, f'view{view_idx:02d}_scale.json')
+            if os.path.exists(src_scale_path):
+                shutil.copy2(src_scale_path, dst_scale_path)
+            
+            records.append({
+                'sha256': sha256,
+                f'shape_latent_view{view_idx:02d}_encoded': True,
+                f'shape_latent_view{view_idx:02d}_tokens': pack['coords'].shape[0]
+            })
+            
+        for _ in tqdm(range(len(tasks)), desc="Extracting view latents"):
+            try:
+                sha256, view_idx, vertices, intersected = load_queue.get()
+                if vertices is None or intersected is None:
+                    continue
+                
+                num_voxels = vertices.feats.shape[0]
+
+                # NaN/Inf check
+                if not (is_valid_sparse_tensor(vertices) and is_valid_sparse_tensor(intersected)):
+                    print(f"[Skip] {sha256}/view{view_idx:02d}: NaN/Inf in input")
+                    continue
+
+                z = encoder(vertices.cuda(), intersected.cuda())
+                torch.cuda.synchronize()
+
+                if not torch.isfinite(z.feats).all():
+                    print(f"[Skip] {sha256}/view{view_idx:02d}: Non-finite latent in z.feats")
+                    clear_cuda_error()
+                    continue
+
+                pack = {
+                    'feats': z.feats.cpu().numpy().astype(np.float32),
+                    'coords': z.coords[:, 1:].cpu().numpy().astype(np.uint8),
+                }
+                saver_executor.submit(saver, sha256, view_idx, pack)
+
+            except Exception as e:
+                print(f"[Error] {sha256}/view{view_idx:02d} ({num_voxels} voxels): {e}")
+                clear_cuda_error()
+                continue
+            
+        saver_executor.shutdown(wait=True)
+        
+    records = pd.DataFrame.from_records(records)
+    records.to_csv(os.path.join(opt.shape_latent_root, 'shape_latents', latent_view_name, 'new_records', f'part_{opt.rank}.csv'), index=False)
diff --git a/data_toolkit/encode_ss_latent.py b/data_toolkit/encode_ss_latent.py
new file mode 100644
index 0000000..ed9f232
--- /dev/null
+++ b/data_toolkit/encode_ss_latent.py
@@ -0,0 +1,163 @@
+import os
+import sys
+sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
+import json
+import argparse
+import torch
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+from easydict import EasyDict as edict
+from concurrent.futures import ThreadPoolExecutor
+from queue import Queue
+
+import pixal3d.models as models
+
+torch.set_grad_enabled(False)
+
+def is_valid_sparse_tensor(tensor):
+    return torch.isfinite(tensor.feats).all() and torch.isfinite(tensor.coords).all()
+
+def clear_cuda_error():
+    torch.cuda.synchronize()
+    torch.cuda.empty_cache()
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--root', type=str, required=True,
+                        help='Directory to save the metadata')
+    parser.add_argument('--shape_latent_root', type=str, default=None,
+                        help='Directory to save the shape latent files')
+    parser.add_argument('--ss_latent_root', type=str, default=None,
+                        help='Directory to save the shape latent files')
+    parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
+                        help='Filter objects with aesthetic score lower than this value')
+    parser.add_argument('--resolution', type=int, default=64,
+                        help='Sparse voxel resolution')
+    parser.add_argument('--shape_latent_name', type=str, default=None,
+                        help='Name of the shape latent files')
+    parser.add_argument('--enc_pretrained', type=str, default='microsoft/TRELLIS-image-large/ckpts/ss_enc_conv3d_16l8_fp16',
+                        help='Pretrained encoder model')
+    parser.add_argument('--model_root', type=str,
+                        help='Root directory of models')
+    parser.add_argument('--enc_model', type=str,
+                        help='Encoder model. if specified, use this model instead of pretrained model')
+    parser.add_argument('--ckpt', type=str,
+                        help='Checkpoint to load')
+    parser.add_argument('--instances', type=str, default=None,
+                        help='Instances to process')
+    parser.add_argument('--rank', type=int, default=0)
+    parser.add_argument('--world_size', type=int, default=1)
+    opt = parser.parse_args()
+    opt = edict(vars(opt))
+    opt.shape_latent_root = opt.shape_latent_root or opt.root
+    opt.ss_latent_root = opt.ss_latent_root or opt.root
+
+    if opt.enc_model is None:
+        latent_name = f'{opt.enc_pretrained.split("/")[-1]}_{opt.resolution}'
+        encoder = models.from_pretrained(opt.enc_pretrained).eval().cuda()
+    else:
+        latent_name = f'{opt.enc_model.split("/")[-1]}_{opt.ckpt}_{opt.resolution}'
+        cfg = edict(json.load(open(os.path.join(opt.model_root, opt.enc_model, 'config.json'), 'r')))
+        encoder = getattr(models, cfg.models.encoder.name)(**cfg.models.encoder.args).cuda()
+        ckpt_path = os.path.join(opt.model_root, opt.enc_model, 'ckpts', f'encoder_{opt.ckpt}.pt')
+        encoder.load_state_dict(torch.load(ckpt_path), strict=False)
+        encoder.eval()
+        print(f'Loaded model from {ckpt_path}')
+    
+    os.makedirs(os.path.join(opt.ss_latent_root, 'ss_latents', latent_name, 'new_records'), exist_ok=True)
+    
+    # get file list
+    if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
+        raise ValueError('metadata.csv not found')
+    metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
+    if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256'))
+    if os.path.exists(os.path.join(opt.shape_latent_root, 'shape_latents', opt.shape_latent_name, 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.shape_latent_root, 'shape_latents', opt.shape_latent_name,'metadata.csv')).set_index('sha256'))
+    if os.path.exists(os.path.join(opt.ss_latent_root,'ss_latents', latent_name, 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.ss_latent_root,'ss_latents', latent_name,'metadata.csv')).set_index('sha256'))
+    metadata = metadata.reset_index()
+    if opt.instances is None:
+        if opt.filter_low_aesthetic_score is not None:
+            metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
+        metadata = metadata[metadata['shape_latent_encoded'] == True]
+        if 'ss_latent_encoded' in metadata.columns:
+            metadata = metadata[metadata['ss_latent_encoded'] != True]
+    else:
+        if os.path.exists(opt.instances):
+            with open(opt.instances, 'r') as f:
+                instances = f.read().splitlines()
+        else:
+            instances = opt.instances.split(',')
+        metadata = metadata[metadata['sha256'].isin(instances)]
+
+    start = len(metadata) * opt.rank // opt.world_size
+    end = len(metadata) * (opt.rank + 1) // opt.world_size
+    metadata = metadata[start:end]
+    records = []
+    
+    # filter out objects that are already processed
+    sha256_list = os.listdir(os.path.join(opt.ss_latent_root, 'ss_latents'))
+    sha256_list = [os.path.splitext(f)[0] for f in sha256_list if f.endswith('.npz')]
+    for sha256 in sha256_list:
+        records.append({'sha256': sha256, 'ss_latent_encoded': True})
+    print(f'Found {len(sha256_list)} processed objects')
+    metadata = metadata[~metadata['sha256'].isin(sha256_list)]
+        
+    print(f'Processing {len(metadata)} objects...')
+    
+    sha256s = list(metadata['sha256'].values)
+    load_queue = Queue(maxsize=32)
+    with ThreadPoolExecutor(max_workers=32) as loader_executor, \
+         ThreadPoolExecutor(max_workers=32) as saver_executor:
+
+        def loader(sha256):
+            try:
+                coords = np.load(os.path.join(opt.shape_latent_root, 'shape_latents', opt.shape_latent_name, f'{sha256}.npz'))['coords']
+                assert np.all(coords < opt.resolution), f"{sha256}: Invalid coords"
+                coords = torch.from_numpy(coords).long()
+                ss = torch.zeros(1, opt.resolution, opt.resolution, opt.resolution, dtype=torch.long)
+                ss[:, coords[:, 0], coords[:, 1], coords[:, 2]] = 1
+                load_queue.put((sha256, ss))
+            except Exception as e:
+                print(f"[Loader Error] {sha256}: {e}")
+                load_queue.put((sha256, None))
+
+        loader_executor.map(loader, sha256s)
+        
+        def saver(sha256, pack):
+            save_path = os.path.join(opt.ss_latent_root, 'ss_latents', latent_name, f'{sha256}.npz')
+            np.savez_compressed(save_path, **pack)
+            records.append({'sha256': sha256, 'ss_latent_encoded': True})
+            
+        for _ in tqdm(range(len(sha256s)), desc="Extracting latents"):
+            try:
+                sha256, ss = load_queue.get()
+                if ss is None:
+                    print(f"[Skip] {sha256}: Failed to load input")
+                    continue
+                
+                ss = ss.cuda()[None].float()
+                z = encoder(ss, sample_posterior=False)
+                torch.cuda.synchronize()
+
+                if not torch.isfinite(z).all():
+                    print(f"[Skip] {sha256}: Non-finite latent")
+                    clear_cuda_error()
+                    continue
+
+                pack = {
+                    'z': z[0].cpu().numpy(),
+                }
+                saver_executor.submit(saver, sha256, pack)
+
+            except Exception as e:
+                print(f"[Error] {sha256}: {e}")
+                clear_cuda_error()
+                continue
+            
+        saver_executor.shutdown(wait=True)
+        
+    records = pd.DataFrame.from_records(records)
+    records.to_csv(os.path.join(opt.ss_latent_root, 'ss_latents', latent_name, 'new_records', f'part_{opt.rank}.csv'), index=False)
diff --git a/data_toolkit/encode_ss_latent_view.py b/data_toolkit/encode_ss_latent_view.py
new file mode 100644
index 0000000..18d0601
--- /dev/null
+++ b/data_toolkit/encode_ss_latent_view.py
@@ -0,0 +1,257 @@
+import os
+import sys
+sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
+import json
+import shutil
+import argparse
+import torch
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+from easydict import EasyDict as edict
+from concurrent.futures import ThreadPoolExecutor
+from queue import Queue
+from utils import parse_view_indices
+
+import pixal3d.models as models
+
+torch.set_grad_enabled(False)
+
+def clear_cuda_error():
+    torch.cuda.synchronize()
+    torch.cuda.empty_cache()
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--root', type=str, required=True,
+                        help='Directory to save the metadata')
+    parser.add_argument('--shape_latent_root', type=str, default=None,
+                        help='Directory containing the shape latent files')
+    parser.add_argument('--ss_latent_root', type=str, default=None,
+                        help='Directory to save the ss latent files')
+    parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
+                        help='Filter objects with aesthetic score lower than this value')
+    parser.add_argument('--resolution', type=int, default=32,
+                        help='SS latent resolution')
+    parser.add_argument('--shape_latent_name', type=str, required=True,
+                        help='Name of the shape latent files (e.g., shape_enc_next_dc_f16c32_fp16_512)')
+    parser.add_argument('--enc_pretrained', type=str, default='microsoft/TRELLIS-image-large/ckpts/ss_enc_conv3d_16l8_fp16',
+                        help='Pretrained encoder model')
+    parser.add_argument('--model_root', type=str,
+                        help='Root directory of models')
+    parser.add_argument('--enc_model', type=str,
+                        help='Encoder model. if specified, use this model instead of pretrained model')
+    parser.add_argument('--ckpt', type=str,
+                        help='Checkpoint to load')
+    parser.add_argument('--instances', type=str, default=None,
+                        help='Instances to process')
+    parser.add_argument('--view_indices', type=str, default=None,
+                        help='View indices to process, e.g., "0,1,2" or "0-5". None for all views')
+    parser.add_argument('--num_views', type=int, default=24,
+                        help='Total number of views (used when view_indices is None)')
+    parser.add_argument('--rank', type=int, default=0)
+    parser.add_argument('--world_size', type=int, default=1)
+    opt = parser.parse_args()
+    opt = edict(vars(opt))
+    opt.shape_latent_root = opt.shape_latent_root or opt.root
+    opt.ss_latent_root = opt.ss_latent_root or opt.root
+
+    # Parse view_indices
+    view_indices = parse_view_indices(opt.view_indices)
+    if view_indices is None:
+        view_indices = list(range(opt.num_views))
+    
+    print(f'View indices to process: {view_indices}')
+
+    if opt.enc_model is None:
+        latent_name = f'{opt.enc_pretrained.split("/")[-1]}_{opt.resolution}'
+        encoder = models.from_pretrained(opt.enc_pretrained).eval().cuda()
+    else:
+        latent_name = f'{opt.enc_model.split("/")[-1]}_{opt.ckpt}_{opt.resolution}'
+        cfg = edict(json.load(open(os.path.join(opt.model_root, opt.enc_model, 'config.json'), 'r')))
+        encoder = getattr(models, cfg.models.encoder.name)(**cfg.models.encoder.args).cuda()
+        ckpt_path = os.path.join(opt.model_root, opt.enc_model, 'ckpts', f'encoder_{opt.ckpt}.pt')
+        encoder.load_state_dict(torch.load(ckpt_path), strict=False)
+        encoder.eval()
+        print(f'Loaded model from {ckpt_path}')
+    
+    # Multi-view shape_latent and ss_latent directory names
+    shape_latent_view_name = f'{opt.shape_latent_name}_view'
+    ss_latent_view_name = f'{latent_name}_view'
+    
+    os.makedirs(os.path.join(opt.ss_latent_root, 'ss_latents', ss_latent_view_name, 'new_records'), exist_ok=True)
+    
+    # Get file list
+    if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
+        raise ValueError('metadata.csv not found')
+    metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
+    if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')):
+        aesthetic_metadata = pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256')
+        metadata = metadata.join(aesthetic_metadata, how='left', rsuffix='_aesthetic')
+    
+    # Check shape_latent_view metadata
+    shape_latent_view_metadata_path = os.path.join(opt.shape_latent_root, 'shape_latents', shape_latent_view_name, 'metadata.csv')
+    if os.path.exists(shape_latent_view_metadata_path):
+        shape_latent_view_metadata = pd.read_csv(shape_latent_view_metadata_path).set_index('sha256')
+        metadata = metadata.join(shape_latent_view_metadata, how='left', rsuffix='_shape_latent_view')
+        print(f'Loaded shape_latent_view metadata with {len(shape_latent_view_metadata)} records')
+    else:
+        print(f'Warning: shape_latent_view metadata not found at {shape_latent_view_metadata_path}')
+    
+    # Check ss_latent_view metadata (used to skip already completed tasks)
+    ss_latent_view_metadata_path = os.path.join(opt.ss_latent_root, 'ss_latents', ss_latent_view_name, 'metadata.csv')
+    if os.path.exists(ss_latent_view_metadata_path):
+        ss_latent_view_metadata = pd.read_csv(ss_latent_view_metadata_path).set_index('sha256')
+        metadata = metadata.join(ss_latent_view_metadata, how='left', rsuffix='_ss_latent_view')
+        print(f'Loaded ss_latent_view metadata with {len(ss_latent_view_metadata)} records')
+    else:
+        print(f'Warning: ss_latent_view metadata not found at {ss_latent_view_metadata_path}')
+    
+    metadata = metadata.reset_index()
+    
+    if opt.instances is None:
+        if opt.filter_low_aesthetic_score is not None:
+            metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
+        
+        # Filter to objects that have shape_latent_view data
+        # Use first view as indicator
+        first_view_col = f'shape_latent_view{view_indices[0]:02d}_encoded'
+        if first_view_col in metadata.columns:
+            metadata = metadata[metadata[first_view_col] == True]
+        else:
+            print(f'Warning: Column {first_view_col} not found in metadata, will check files directly')
+    else:
+        if os.path.exists(opt.instances):
+            with open(opt.instances, 'r') as f:
+                instances = f.read().splitlines()
+        else:
+            instances = opt.instances.split(',')
+        metadata = metadata[metadata['sha256'].isin(instances)]
+
+    start = len(metadata) * opt.rank // opt.world_size
+    end = len(metadata) * (opt.rank + 1) // opt.world_size
+    metadata = metadata[start:end]
+    records = []
+    
+    # Build task list: (sha256, view_idx), filter already completed tasks via metadata
+    tasks = []
+    skipped_count = 0
+    
+    # Pre-fetch completion status columns for each view
+    encoded_cols = {view_idx: f'ss_latent_view{view_idx:02d}_encoded' for view_idx in view_indices}
+    
+    for _, row in metadata.iterrows():
+        sha256 = row['sha256']
+        for view_idx in view_indices:
+            encoded_col = encoded_cols[view_idx]
+            # Check if already marked as completed in metadata
+            if encoded_col in metadata.columns and row.get(encoded_col, False) == True:
+                skipped_count += 1
+                continue
+            tasks.append((sha256, view_idx))
+    
+    print(f'Total tasks: {len(tasks) + skipped_count}, Already done (from metadata): {skipped_count}, To process: {len(tasks)}')
+    
+    load_queue = Queue(maxsize=32)
+    
+    with ThreadPoolExecutor(max_workers=32) as loader_executor, \
+         ThreadPoolExecutor(max_workers=32) as saver_executor:
+
+        def loader(task):
+            sha256, view_idx = task
+            try:
+                # Check if output file already exists, skip if so
+                output_path = os.path.join(
+                    opt.ss_latent_root,
+                    'ss_latents',
+                    ss_latent_view_name,
+                    sha256,
+                    f'view{view_idx:02d}.npz'
+                )
+                if os.path.exists(output_path):
+                    load_queue.put((sha256, view_idx, None))
+                    return
+                
+                # shape_latent_view path: shape_latents/{shape_latent_view_name}/{sha256}/view{idx:02d}.npz
+                npz_path = os.path.join(
+                    opt.shape_latent_root, 
+                    'shape_latents',
+                    shape_latent_view_name, 
+                    sha256, 
+                    f'view{view_idx:02d}.npz'
+                )
+                
+                if not os.path.exists(npz_path):
+                    print(f"[Loader Skip] {sha256}/view{view_idx:02d}: npz file not found at {npz_path}")
+                    load_queue.put((sha256, view_idx, None))
+                    return
+                
+                data = np.load(npz_path)
+                coords = data['coords']
+                
+                # Validate coords are within resolution range
+                assert np.all(coords < opt.resolution), f"{sha256}/view{view_idx:02d}: Invalid coords (max={coords.max()}, resolution={opt.resolution})"
+                
+                coords = torch.from_numpy(coords).long()
+                ss = torch.zeros(1, opt.resolution, opt.resolution, opt.resolution, dtype=torch.long)
+                ss[:, coords[:, 0], coords[:, 1], coords[:, 2]] = 1
+                
+                load_queue.put((sha256, view_idx, ss))
+            except Exception as e:
+                print(f"[Loader Error] {sha256}/view{view_idx:02d}: {e}")
+                load_queue.put((sha256, view_idx, None))
+
+        loader_executor.map(loader, tasks)
+        
+        def saver(sha256, view_idx, pack):
+            sha256_dir = os.path.join(opt.ss_latent_root, 'ss_latents', ss_latent_view_name, sha256)
+            os.makedirs(sha256_dir, exist_ok=True)
+            save_path = os.path.join(sha256_dir, f'view{view_idx:02d}.npz')
+            np.savez_compressed(save_path, **pack)
+            
+            # Copy scale.json from shape_latent_view directory
+            src_scale_path = os.path.join(
+                opt.shape_latent_root,
+                'shape_latents',
+                shape_latent_view_name,
+                sha256,
+                f'view{view_idx:02d}_scale.json'
+            )
+            dst_scale_path = os.path.join(sha256_dir, f'view{view_idx:02d}_scale.json')
+            if os.path.exists(src_scale_path) and not os.path.exists(dst_scale_path):
+                shutil.copy2(src_scale_path, dst_scale_path)
+            
+            records.append({
+                'sha256': sha256,
+                f'ss_latent_view{view_idx:02d}_encoded': True,
+            })
+            
+        for _ in tqdm(range(len(tasks)), desc="Extracting SS view latents"):
+            try:
+                sha256, view_idx, ss = load_queue.get()
+                if ss is None:
+                    continue
+                
+                ss = ss.cuda()[None].float()
+                z = encoder(ss, sample_posterior=False)
+                torch.cuda.synchronize()
+
+                if not torch.isfinite(z).all():
+                    print(f"[Skip] {sha256}/view{view_idx:02d}: Non-finite latent")
+                    clear_cuda_error()
+                    continue
+
+                pack = {
+                    'z': z[0].cpu().numpy(),
+                }
+                saver_executor.submit(saver, sha256, view_idx, pack)
+
+            except Exception as e:
+                print(f"[Error] {sha256}/view{view_idx:02d}: {e}")
+                clear_cuda_error()
+                continue
+            
+        saver_executor.shutdown(wait=True)
+        
+    records = pd.DataFrame.from_records(records)
+    records.to_csv(os.path.join(opt.ss_latent_root, 'ss_latents', ss_latent_view_name, 'new_records', f'part_{opt.rank}.csv'), index=False)
diff --git a/data_toolkit/render_cond.py b/data_toolkit/render_cond.py
new file mode 100644
index 0000000..14067ff
--- /dev/null
+++ b/data_toolkit/render_cond.py
@@ -0,0 +1,154 @@
+import os
+import json
+import copy
+import sys
+import importlib
+import argparse
+import pandas as pd
+from easydict import EasyDict as edict
+from functools import partial
+from subprocess import DEVNULL, call
+from concurrent.futures import ThreadPoolExecutor
+from tqdm import tqdm
+import numpy as np
+from utils import sphere_hammersley_sequence
+
+
+BLENDER_LINK = 'https://download.blender.org/release/Blender3.0/blender-3.0.1-linux-x64.tar.xz'
+BLENDER_INSTALLATION_PATH = '/tmp'
+BLENDER_PATH = f'{BLENDER_INSTALLATION_PATH}/blender-3.0.1-linux-x64/blender'
+
+def _install_blender():
+    if not os.path.exists(BLENDER_PATH):
+        os.system('sudo apt-get update')
+        os.system('sudo apt-get install -y libxrender1 libxi6 libxkbcommon-x11-0 libsm6 libxfixes3 libgl1')
+        os.system(f'wget {BLENDER_LINK} -P {BLENDER_INSTALLATION_PATH}')
+        os.system(f'tar -xvf {BLENDER_INSTALLATION_PATH}/blender-3.0.1-linux-x64.tar.xz -C {BLENDER_INSTALLATION_PATH}')
+    # Install Pillow into Blender's bundled Python if not already present
+    blender_python = os.path.join(BLENDER_INSTALLATION_PATH, 'blender-3.0.1-linux-x64', '3.0', 'python', 'bin', 'python3.9')
+    try:
+        import subprocess
+        subprocess.check_call([blender_python, '-c', 'from PIL import Image'], stdout=DEVNULL, stderr=DEVNULL)
+    except subprocess.CalledProcessError:
+        print('Installing Pillow into Blender Python...')
+        subprocess.check_call([blender_python, '-m', 'ensurepip'], stdout=DEVNULL, stderr=DEVNULL)
+        subprocess.check_call([blender_python, '-m', 'pip', 'install', 'Pillow'], stdout=DEVNULL, stderr=DEVNULL)
+
+
+def _render_cond(file_path, sha256, root, num_cond_views):
+    # Build conditional view camera
+    yaws = []
+    pitchs = []
+    offset = (np.random.rand(), np.random.rand())
+    for i in range(num_cond_views):
+        y, p = sphere_hammersley_sequence(i, num_cond_views, offset)
+        yaws.append(y)
+        pitchs.append(p)
+    fov_min, fov_max = 10, 70
+    radius_min = np.sqrt(3) / 2 / np.sin(fov_max / 360 * np.pi)
+    radius_max = np.sqrt(3) / 2 / np.sin(fov_min / 360 * np.pi)
+    k_min = 1 / radius_max**2
+    k_max = 1 / radius_min**2
+    ks = np.random.uniform(k_min, k_max, (1000000,))
+    radius = [1 / np.sqrt(k) for k in ks]
+    fov = [2 * np.arcsin(np.sqrt(3) / 2 / r) for r in radius]
+    cond_views = [{'yaw': y, 'pitch': p, 'radius': r, 'fov': f} for y, p, r, f in zip(yaws, pitchs, radius, fov)]
+    
+    args = [
+        BLENDER_PATH, '-b', '-P', os.path.join(os.path.dirname(__file__), 'blender_script', 'render_cond.py'),
+        '--',
+        '--object', os.path.expanduser(file_path),
+        '--cond_views', json.dumps(cond_views),
+        '--cond_resolution', '1024',
+        '--cond_output_folder', os.path.join(root, 'renders_cond', sha256),
+        '--engine', 'CYCLES',
+    ]
+    if file_path.endswith('.blend'):
+        args.insert(1, file_path)
+    
+    call(args, stdout=DEVNULL, stderr=DEVNULL)
+    
+    if os.path.exists(os.path.join(root, 'renders_cond', sha256, 'transforms.json')):
+        return {'sha256': sha256, 'cond_rendered': True}
+
+
+if __name__ == '__main__':
+    dataset_utils = importlib.import_module(f'datasets.{sys.argv[1]}')
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--root', type=str, required=True,
+                        help='Directory to save the metadata')
+    parser.add_argument('--download_root', type=str, default=None,
+                        help='Directory to save the downloaded files')
+    parser.add_argument('--render_cond_root', type=str, default=None,
+                        help='Directory to save the mesh dumps')
+    parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
+                        help='Filter objects with aesthetic score lower than this value')
+    parser.add_argument('--instances', type=str, default=None,
+                        help='Instances to process')
+    parser.add_argument('--num_cond_views', type=int, default=2,
+                        help='Number of conditional views to render')
+    dataset_utils.add_args(parser)
+    parser.add_argument('--rank', type=int, default=0)
+    parser.add_argument('--world_size', type=int, default=1)
+    parser.add_argument('--max_workers', type=int, default=8)
+    opt = parser.parse_args(sys.argv[2:])
+    opt = edict(vars(opt))
+    opt.download_root = opt.download_root or opt.root
+    opt.render_cond_root = opt.render_cond_root or opt.root
+
+    os.makedirs(os.path.join(opt.render_cond_root, 'renders_cond', 'new_records'), exist_ok=True)
+    
+    # install blender
+    print('Checking blender...', flush=True)
+    _install_blender()
+
+    # get file list
+    if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
+        raise ValueError('metadata.csv not found')
+    metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
+    if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256'))
+    if os.path.exists(os.path.join(opt.download_root, 'raw', 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.download_root, 'raw', 'metadata.csv')).set_index('sha256'))
+    if os.path.exists(os.path.join(opt.render_cond_root, 'renders_cond', 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.render_cond_root, 'renders_cond', 'metadata.csv')).set_index('sha256'))
+    metadata = metadata.reset_index()
+    if opt.instances is None:
+        metadata = metadata[metadata['local_path'].notna()]
+        if opt.filter_low_aesthetic_score is not None:
+            metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
+        if 'cond_rendered' in metadata.columns:
+            metadata = metadata[(metadata['cond_rendered'] != True)]
+    else:
+        if os.path.exists(opt.instances):
+            with open(opt.instances, 'r') as f:
+                instances = f.read().splitlines()
+        else:
+            instances = opt.instances.split(',')
+        metadata = metadata[metadata['sha256'].isin(instances)]
+
+    start = len(metadata) * opt.rank // opt.world_size
+    end = len(metadata) * (opt.rank + 1) // opt.world_size
+    metadata = metadata[start:end]
+    records = []
+
+    # filter out objects that are already processed
+    with ThreadPoolExecutor(max_workers=os.cpu_count()) as executor, \
+        tqdm(total=len(metadata), desc="Filtering existing objects") as pbar:
+        def check_sha256(sha256):
+            if os.path.exists(os.path.join(opt.render_cond_root, 'renders_cond', sha256, 'transforms.json')):
+                records.append({'sha256': sha256, 'cond_rendered': True})
+            pbar.update()
+        executor.map(check_sha256, metadata['sha256'].values)
+        executor.shutdown(wait=True)
+    existing_sha256 = set(r['sha256'] for r in records)
+    metadata = metadata[~metadata['sha256'].isin(existing_sha256)]
+
+    print(f'Processing {len(metadata)} objects...')
+
+    # process objects
+    func = partial(_render_cond, root=opt.render_cond_root, num_cond_views=opt.num_cond_views)
+    cond_rendered = dataset_utils.foreach_instance(metadata, opt.render_cond_root, func, max_workers=opt.max_workers, desc='Rendering objects')
+    cond_rendered = pd.concat([cond_rendered, pd.DataFrame.from_records(records)])
+    cond_rendered.to_csv(os.path.join(opt.render_cond_root, 'renders_cond', 'new_records', f'part_{opt.rank}.csv'), index=False)
diff --git a/data_toolkit/setup.sh b/data_toolkit/setup.sh
new file mode 100644
index 0000000..3cb1427
--- /dev/null
+++ b/data_toolkit/setup.sh
@@ -0,0 +1 @@
+pip install pillow imageio imageio-ffmpeg tqdm easydict opencv-python-headless pandas open3d objaverse huggingface_hub[cli] open_clip_torch
diff --git a/data_toolkit/utils.py b/data_toolkit/utils.py
new file mode 100644
index 0000000..c9ede78
--- /dev/null
+++ b/data_toolkit/utils.py
@@ -0,0 +1,564 @@
+from typing import *
+import hashlib
+import numpy as np
+import cv2
+
+
+def get_file_hash(file: str) -> str:
+    sha256 = hashlib.sha256()
+    # Read the file from the path
+    with open(file, "rb") as f:
+        # Update the hash with the file content
+        for byte_block in iter(lambda: f.read(4096), b""):
+            sha256.update(byte_block)
+    return sha256.hexdigest()
+
+# ===============LOW DISCREPANCY SEQUENCES================
+
+PRIMES = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53]
+
+def radical_inverse(base, n):
+    val = 0
+    inv_base = 1.0 / base
+    inv_base_n = inv_base
+    while n > 0:
+        digit = n % base
+        val += digit * inv_base_n
+        n //= base
+        inv_base_n *= inv_base
+    return val
+
+def halton_sequence(dim, n):
+    return [radical_inverse(PRIMES[dim], n) for dim in range(dim)]
+
+def hammersley_sequence(dim, n, num_samples):
+    return [n / num_samples] + halton_sequence(dim - 1, n)
+
+def sphere_hammersley_sequence(n, num_samples, offset=(0, 0)):
+    u, v = hammersley_sequence(2, n, num_samples)
+    u += offset[0] / num_samples
+    v += offset[1]
+    u = 2 * u if u < 0.25 else 2 / 3 * u + 1 / 3
+    theta = np.arccos(1 - 2 * u) - np.pi / 2
+    phi = v * 2 * np.pi
+    return [phi, theta]
+
+# ==============PLY IO===============
+import struct
+import re
+import torch
+
+def read_ply(filename):
+    """
+    Read a PLY file and return vertices, triangle faces, and quad faces.
+    
+    Args:
+        filename (str): The file path to read from.
+        
+    Returns:
+        vertices (torch.Tensor): Tensor of shape [N, 3] containing vertex positions.
+        tris (torch.Tensor): Tensor of shape [M, 3] containing triangle face indices (empty if none).
+        quads (torch.Tensor): Tensor of shape [K, 4] containing quad face indices (empty if none).
+    """
+    with open(filename, 'rb') as f:
+        # Read the header until 'end_header' is encountered
+        header_bytes = b""
+        while True:
+            line = f.readline()
+            if not line:
+                raise ValueError("PLY header not found")
+            header_bytes += line
+            if b"end_header" in line:
+                break
+        header = header_bytes.decode('utf-8')
+        
+        # Determine if the file is in ASCII or binary format
+        is_ascii = "ascii" in header
+        
+        # Extract the number of vertices and faces from the header using regex
+        vertex_match = re.search(r'element vertex (\d+)', header)
+        if vertex_match:
+            num_vertices = int(vertex_match.group(1))
+        else:
+            raise ValueError("Vertex count not found in header")
+            
+        face_match = re.search(r'element face (\d+)', header)
+        if face_match:
+            num_faces = int(face_match.group(1))
+        else:
+            raise ValueError("Face count not found in header")
+        
+        vertices = []
+        tris = []
+        quads = []
+        
+        if is_ascii:
+            # For ASCII format, read each line of vertex data (each line contains 3 floats)
+            for _ in range(num_vertices):
+                line = f.readline().decode('utf-8').strip()
+                if not line: 
+                    continue
+                parts = line.split()
+                vertices.append([float(parts[0]), float(parts[1]), float(parts[2])])
+            
+            # Read face data, where the first number indicates the number of vertices for the face
+            for _ in range(num_faces):
+                line = f.readline().decode('utf-8').strip()
+                if not line: 
+                    continue
+                parts = line.split()
+                count = int(parts[0])
+                indices = list(map(int, parts[1:]))
+                if count == 3:
+                    tris.append(indices)
+                elif count == 4:
+                    quads.append(indices)
+                else:
+                    # Skip faces with other numbers of vertices (can be extended as needed)
+                    pass
+        else:
+            # For binary format: read directly from the binary stream
+            # Each vertex consists of 3 floats (12 bytes per vertex)
+            for _ in range(num_vertices):
+                data = f.read(12)
+                if len(data) < 12:
+                    raise ValueError("Insufficient vertex data")
+                v = struct.unpack('<fff', data)
+                vertices.append(v)
+            
+            # Read face data from the binary stream
+            for _ in range(num_faces):
+                # First, read 1 byte indicating the number of vertices in the face
+                count_data = f.read(1)
+                if len(count_data) < 1:
+                    raise ValueError("Failed to read face vertex count")
+                count = struct.unpack('<B', count_data)[0]
+                if count == 3:
+                    data = f.read(12)  # 3 * 4 bytes
+                    if len(data) < 12:
+                        raise ValueError("Insufficient data for triangle face")
+                    indices = struct.unpack('<3i', data)
+                    tris.append(indices)
+                elif count == 4:
+                    data = f.read(16)  # 4 * 4 bytes
+                    if len(data) < 16:
+                        raise ValueError("Insufficient data for quad face")
+                    indices = struct.unpack('<4i', data)
+                    quads.append(indices)
+                else:
+                    # For faces with a different number of vertices, read count*4 bytes
+                    data = f.read(count * 4)
+                    # Skip or extend processing as needed
+                    raise ValueError(f"Unsupported face with {count} vertices")
+        
+        # Convert lists to torch.Tensor
+        vertices = torch.tensor(vertices, dtype=torch.float32)
+        tris = torch.tensor(tris, dtype=torch.int32) if len(tris) > 0 else torch.empty((0, 3), dtype=torch.int32)
+        quads = torch.tensor(quads, dtype=torch.int32) if len(quads) > 0 else torch.empty((0, 4), dtype=torch.int32)
+        
+        return vertices, tris, quads
+
+
+def write_ply(filename, vertices, tris, quads, ascii=False):
+    """
+    Write a mesh to a PLY file, with the option to save in ASCII or binary format.
+    
+    Args:
+        filename (str): The filename to write to.
+        vertices (torch.Tensor): [N, 3] The vertex positions.
+        tris (torch.Tensor): [M, 3] The triangle indices.
+        quads (torch.Tensor): [K, 4] The quad indices.
+        ascii (bool): If True, write in ASCII format. If False, write in binary format.
+    """
+    # Convert torch tensors to numpy arrays
+    vertices = vertices.numpy()
+    tris = tris.numpy()
+    quads = quads.numpy()
+
+    # Prepare the header
+    num_vertices = len(vertices)
+    num_faces = len(tris) + len(quads)
+
+    # Vertex properties
+    vertex_header = "property float x\nproperty float y\nproperty float z"
+
+    # Face properties (the number of vertices per face is variable)
+    face_header = "property list uchar int vertex_index"
+
+    # Start writing the PLY header
+    header = f"ply\n"
+    header += f"format {'ascii 1.0' if ascii else 'binary_little_endian 1.0'}\n"
+    header += f"element vertex {num_vertices}\n"
+    header += vertex_header + "\n"
+    header += f"element face {num_faces}\n"
+    header += face_header + "\n"
+    header += "end_header\n"
+
+    # Open the file for writing
+    with open(filename, 'wb' if not ascii else 'w') as f:
+        # Write the header
+        f.write(header if ascii else header.encode('utf-8'))
+
+        # Write the vertex data
+        if ascii:
+            for v in vertices:
+                f.write(f"{v[0]} {v[1]} {v[2]}\n")
+        else:
+            for v in vertices:
+                f.write(struct.pack('<fff', *v))
+
+        # Write the face data
+        if ascii:
+            for tri in tris:
+                f.write(f"3 {tri[0]} {tri[1]} {tri[2]}\n")
+            for quad in quads:
+                f.write(f"4 {quad[0]} {quad[1]} {quad[2]} {quad[3]}\n")
+        else:
+            for tri in tris:
+                f.write(struct.pack('<B3i', 3, *tri))  # 3 indices for triangle
+            for quad in quads:
+                f.write(struct.pack('<B4i', 4, *quad))  # 4 indices for quad
+                
+                
+# ==============IMAGE UTILS===============
+
+def make_grid(images, nrow=None, ncol=None, aspect_ratio=None):
+    num_images = len(images)
+    if nrow is None and ncol is None:
+        if aspect_ratio is not None:
+            nrow = int(np.round(np.sqrt(num_images / aspect_ratio)))
+        else:
+            nrow = int(np.sqrt(num_images))
+        ncol = (num_images + nrow - 1) // nrow
+    elif nrow is None and ncol is not None:
+        nrow = (num_images + ncol - 1) // ncol
+    elif nrow is not None and ncol is None:
+        ncol = (num_images + nrow - 1) // nrow
+    else:
+        assert nrow * ncol >= num_images, 'nrow * ncol must be greater than or equal to the number of images'
+    
+    if images[0].ndim == 2:
+        grid = np.zeros((nrow * images[0].shape[0], ncol * images[0].shape[1]), dtype=images[0].dtype)
+    else:
+        grid = np.zeros((nrow * images[0].shape[0], ncol * images[0].shape[1], images[0].shape[2]), dtype=images[0].dtype)
+    for i, img in enumerate(images):
+        row = i // ncol
+        col = i % ncol
+        grid[row * img.shape[0]:(row + 1) * img.shape[0], col * img.shape[1]:(col + 1) * img.shape[1]] = img
+    return grid
+
+
+def notes_on_image(img, notes=None):
+    img = np.pad(img, ((0, 32), (0, 0), (0, 0)), 'constant', constant_values=0)
+    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+    if notes is not None:
+        img = cv2.putText(img, notes, (0, img.shape[0] - 4), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1)
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    return img
+
+
+
+def text_image(text, resolution=(512, 512), max_size=0.5, h_align="left", v_align="center"):
+    """
+    Draw text on an image of the given resolution. The text is automatically wrapped
+    and scaled so that it fits completely within the image while preserving any explicit
+    line breaks and original spacing. Horizontal and vertical alignment can be controlled
+    via flags.
+    
+    Parameters:
+        text (str): The input text. Newline characters and spacing are preserved.
+        resolution (tuple): The image resolution as (width, height).
+        max_size (float): The maximum font size.
+        h_align (str): Horizontal alignment. Options: "left", "center", "right".
+        v_align (str): Vertical alignment. Options: "top", "center", "bottom".
+        
+    Returns:
+        numpy.ndarray: The resulting image (BGR format) with the text drawn.
+    """
+    width, height = resolution
+    # Create a white background image
+    img = np.full((height, width, 3), 255, dtype=np.uint8)
+
+    # Set margins and compute available drawing area
+    margin = 10
+    avail_width = width - 2 * margin
+    avail_height = height - 2 * margin
+
+    # Choose OpenCV font and text thickness
+    font = cv2.FONT_HERSHEY_SIMPLEX
+    thickness = 1
+    # Ratio for additional spacing between lines (relative to the height of "A")
+    line_spacing_ratio = 0.5
+
+    def wrap_line(line, max_width, font, thickness, scale):
+        """
+        Wrap a single line of text into multiple lines such that each line's
+        width (measured at the given scale) does not exceed max_width.
+        This function preserves the original spacing by splitting the line into tokens
+        (words and whitespace) using a regular expression.
+        
+        Parameters:
+            line (str): The input text line.
+            max_width (int): Maximum allowed width in pixels.
+            font (int): OpenCV font identifier.
+            thickness (int): Text thickness.
+            scale (float): The current font scale.
+            
+        Returns:
+            List[str]: A list of wrapped lines.
+        """
+        # Split the line into tokens (words and whitespace), preserving spacing
+        tokens = re.split(r'(\s+)', line)
+        if not tokens:
+            return ['']
+        
+        wrapped_lines = []
+        current_line = ""
+        for token in tokens:
+            candidate = current_line + token
+            candidate_width = cv2.getTextSize(candidate, font, scale, thickness)[0][0]
+            if candidate_width <= max_width:
+                current_line = candidate
+            else:
+                # If current_line is empty, the token itself is too wide;
+                # break the token character by character.
+                if current_line == "":
+                    sub_token = ""
+                    for char in token:
+                        candidate_char = sub_token + char
+                        if cv2.getTextSize(candidate_char, font, scale, thickness)[0][0] <= max_width:
+                            sub_token = candidate_char
+                        else:
+                            if sub_token:
+                                wrapped_lines.append(sub_token)
+                            sub_token = char
+                    current_line = sub_token
+                else:
+                    wrapped_lines.append(current_line)
+                    current_line = token
+        if current_line:
+            wrapped_lines.append(current_line)
+        return wrapped_lines
+
+    def compute_text_block(scale):
+        """
+        Wrap the entire text (splitting at explicit newline characters) using the
+        provided scale, and then compute the overall width and height of the text block.
+        
+        Returns:
+            wrapped_lines (List[str]): The list of wrapped lines.
+            block_width (int): Maximum width among the wrapped lines.
+            block_height (int): Total height of the text block including spacing.
+            sizes (List[tuple]): A list of (width, height) for each wrapped line.
+            spacing (int): The spacing between lines (computed from the scaled "A" height).
+        """
+        # Split text by explicit newlines
+        input_lines = text.splitlines() if text else ['']
+        wrapped_lines = []
+        for line in input_lines:
+            wrapped = wrap_line(line, avail_width, font, thickness, scale)
+            wrapped_lines.extend(wrapped)
+            
+        sizes = []
+        for line in wrapped_lines:
+            (text_size, _) = cv2.getTextSize(line, font, scale, thickness)
+            sizes.append(text_size)  # (width, height)
+            
+        block_width = max((w for w, h in sizes), default=0)
+        # Use the height of "A" (at the current scale) to compute line spacing
+        base_height = cv2.getTextSize("A", font, scale, thickness)[0][1]
+        spacing = int(line_spacing_ratio * base_height)
+        block_height = sum(h for w, h in sizes) + spacing * (len(sizes) - 1) if sizes else 0
+        
+        return wrapped_lines, block_width, block_height, sizes, spacing
+
+    # Use binary search to find the maximum scale that allows the text block to fit
+    lo = 0.001
+    hi = max_size
+    eps = 0.001  # convergence threshold
+    best_scale = lo
+    best_result = None
+
+    while hi - lo > eps:
+        mid = (lo + hi) / 2
+        wrapped_lines, block_width, block_height, sizes, spacing = compute_text_block(mid)
+        # Ensure that both width and height constraints are met
+        if block_width <= avail_width and block_height <= avail_height:
+            best_scale = mid
+            best_result = (wrapped_lines, block_width, block_height, sizes, spacing)
+            lo = mid  # try a larger scale
+        else:
+            hi = mid  # reduce the scale
+
+    if best_result is None:
+        best_scale = 0.5
+        best_result = compute_text_block(best_scale)
+        
+    wrapped_lines, block_width, block_height, sizes, spacing = best_result
+
+    # Compute starting y-coordinate based on vertical alignment flag
+    if v_align == "top":
+        y_top = margin
+    elif v_align == "center":
+        y_top = margin + (avail_height - block_height) // 2
+    elif v_align == "bottom":
+        y_top = margin + (avail_height - block_height)
+    else:
+        y_top = margin + (avail_height - block_height) // 2  # default to center if invalid flag
+
+    # For cv2.putText, the y coordinate represents the text baseline;
+    # so for the first line add its height.
+    y = y_top + (sizes[0][1] if sizes else 0)
+
+    # Draw each line with horizontal alignment based on the flag
+    for i, line in enumerate(wrapped_lines):
+        line_width, line_height = sizes[i]
+        if h_align == "left":
+            x = margin
+        elif h_align == "center":
+            x = margin + (avail_width - line_width) // 2
+        elif h_align == "right":
+            x = margin + (avail_width - line_width)
+        else:
+            x = margin  # default to left if invalid flag
+
+        cv2.putText(img, line, (x, y), font, best_scale, (0, 0, 0), thickness, cv2.LINE_AA)
+        y += line_height + spacing
+
+    return img
+
+
+# ==================== View index parsing ====================
+
+def parse_view_indices(view_indices_str):
+    """Parse view_indices string into a sorted deduplicated list of integers."""
+    if view_indices_str is None:
+        return None
+    view_indices = []
+    for part in view_indices_str.split(','):
+        if '-' in part:
+            start, end = map(int, part.split('-'))
+            view_indices.extend(range(start, end + 1))
+        else:
+            view_indices.append(int(part))
+    view_indices = list(set(view_indices))
+    view_indices.sort()
+    return view_indices
+
+
+# ==================== Multi-view transform functions ====================
+
+import math
+
+def get_new_camera_matrix(radius: float, yaw: float, pitch: float, dtype=torch.float32, device='cpu'):
+    """
+    Compute camera-to-world 4x4 transform matrix in spherical coordinates,
+    looking at origin with up=(0,1,0).
+    Uses standard LookAt formula where camera local -Z points at target and local Y is up.
+    An additional +90 degree rotation around local Z axis is applied to match
+    Blender Track-To local axis convention.
+    yaw, pitch are in radians.
+    """
+    x = radius * math.cos(yaw) * math.cos(pitch)
+    y = radius * math.sin(yaw) * math.cos(pitch)
+    z = radius * math.sin(pitch)
+    eye = torch.tensor([x, y, z], dtype=dtype, device=device)
+    target = torch.zeros(3, dtype=dtype, device=device)
+    up_global = torch.tensor([0.0, 1.0, 0.0], dtype=dtype, device=device)
+
+    f = (target - eye)
+    f = f / torch.norm(f)
+    r = torch.cross(f, up_global)
+    r = r / torch.norm(r)
+    u = torch.cross(r, f)
+
+    z_cam = -f
+    x_cam = r
+    y_cam = u
+
+    T = torch.eye(4, dtype=dtype, device=device)
+    T[:3, 0] = x_cam
+    T[:3, 1] = y_cam
+    T[:3, 2] = z_cam
+    T[:3, 3] = eye
+
+    # +90 degree rotation matrix around local Z axis
+    Rz90 = torch.tensor([
+        [0.0, -1.0, 0.0, 0.0],
+        [1.0,  0.0, 0.0, 0.0],
+        [0.0,  0.0, 1.0, 0.0],
+        [0.0,  0.0, 0.0, 1.0]
+    ], dtype=dtype, device=device)
+
+    return T @ Rz90
+
+
+def transform_mesh(mesh_v, frame):
+    """
+    Apply multi-view transform to mesh vertices based on camera transform matrix.
+    """
+    device = mesh_v.device
+    c2w_orig = torch.tensor(frame['transform_matrix'], dtype=torch.float32, device=device)
+
+    # Old and new camera matrices
+    radius = c2w_orig[:3, 3].norm().item()
+    c2w_new = get_new_camera_matrix(radius=radius, yaw=-90/180.0*math.pi, pitch=0.0,
+                                dtype=torch.float32, device=device)
+    w2c_orig = torch.inverse(c2w_orig)
+
+    # Initial and final axis alignment matrices
+    R_init = torch.tensor([
+        [1.0, 0.0,  0.0, 0.0],
+        [0.0, 0.0, -1.0, 0.0],
+        [0.0, 1.0,  0.0, 0.0],
+        [0.0, 0.0,  0.0, 1.0]
+    ], dtype=torch.float32, device=device)
+
+    R_back = torch.tensor([
+        [1.0,  0.0, 0.0, 0.0],
+        [0.0,  0.0, 1.0, 0.0],
+        [0.0, -1.0, 0.0, 0.0],
+        [0.0,  0.0, 0.0, 1.0]
+    ], dtype=torch.float32, device=device)
+
+    R_ply = torch.tensor([
+        [1.0,  0.0, 0.0, 0.0],
+        [0.0,  0.0, 1.0, 0.0],
+        [0.0, -1.0, 0.0, 0.0],
+        [0.0,  0.0, 0.0, 1.0]
+    ], dtype=torch.float32, device=device)
+
+    T_cam = c2w_new @ w2c_orig @ R_ply
+    T_final = R_back @ T_cam @ R_init
+
+    # Apply transform
+    mesh_v = mesh_v.reshape(-1, 3)
+    verts_h = torch.cat([mesh_v, torch.ones((mesh_v.shape[0], 1), dtype=torch.float32, device=device)], dim=1)
+    verts_trans = (T_final @ verts_h.T).T[:, :3]
+
+    return verts_trans
+
+
+def sphere_normalize_torch(vertices):
+    """
+    Sphere normalization: normalize vertices based on sphere radius.
+    """
+    bmin = torch.min(vertices, dim=0)[0]
+    bmax = torch.max(vertices, dim=0)[0]
+    bcenter = (bmax + bmin) / 2
+    assert bcenter.abs().max() < 0.25, f"bcenter is not close to origin: {bcenter}"
+    radius = torch.norm(vertices - bcenter, dim=-1).max()
+    vertices_normalized = vertices / radius
+    return vertices_normalized, bcenter, radius
+
+
+def save_image_with_notes(img, path, notes=None):
+    """
+    Save an image with notes.
+    """
+    if isinstance(img, torch.Tensor):
+        img = img.cpu().numpy().transpose(1, 2, 0)
+    if img.dtype == np.float32 or img.dtype == np.float64:
+        img = np.clip(img * 255, 0, 255).astype(np.uint8)
+    img = notes_on_image(img, notes)
+    cv2.imwrite(path, cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
diff --git a/data_toolkit/visualize_pbr_latent.py b/data_toolkit/visualize_pbr_latent.py
new file mode 100644
index 0000000..18f9ca1
--- /dev/null
+++ b/data_toolkit/visualize_pbr_latent.py
@@ -0,0 +1,237 @@
+"""
+Decode shape + PBR latent (.npz) to textured GLB mesh and render the PBR front view.
+
+Usage:
+    python data_toolkit/visualize_pbr_latent.py \
+        --root datasets/ObjaverseXL_sketchfab \
+        --sha256 <SHA256_HASH> \
+        --resolution 1024 \
+        --view_idx 0
+"""
+
+import os
+import sys
+import json
+import shutil
+import argparse
+import numpy as np
+import torch
+import cv2
+from PIL import Image
+
+os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1'
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
+
+import pixal3d.models as models
+import pixal3d.modules.sparse as sp
+from pixal3d.representations import MeshWithVoxel
+from pixal3d.renderers import EnvMap
+from pixal3d.utils import render_utils
+import o_voxel
+
+PBR_ATTR_LAYOUT = {
+    'base_color': slice(0, 3),
+    'metallic': slice(3, 4),
+    'roughness': slice(4, 5),
+    'alpha': slice(5, 6),
+}
+
+
+def load_latent(latent_file):
+    """Load a latent .npz file and return a SparseTensor on GPU."""
+    data = np.load(latent_file)
+    coords = torch.tensor(data['coords']).int()
+    feats = torch.tensor(data['feats']).float()
+    coords = torch.cat([torch.zeros_like(coords[:, :1]), coords], dim=1)
+    return sp.SparseTensor(feats.cuda(), coords.cuda())
+
+
+def load_envmaps(device='cuda'):
+    """Load HDRI environment maps from assets/."""
+    base = os.path.join(os.path.dirname(__file__), '..', 'assets', 'hdri')
+    envmaps = {}
+    for name in ['forest', 'sunset', 'courtyard']:
+        path = os.path.join(base, f'{name}.exr')
+        if os.path.exists(path):
+            img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+            envmaps[name] = EnvMap(torch.tensor(img, dtype=torch.float32, device=device))
+    return envmaps
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Decode shape + PBR latent to textured GLB and render")
+    parser.add_argument("--root", type=str, required=True, help="Dataset root")
+    parser.add_argument("--sha256", type=str, required=True, help="SHA256 of the asset")
+    parser.add_argument("--resolution", type=int, default=1024, help="Decoder resolution")
+    parser.add_argument("--view_idx", type=int, default=0, help="View index to decode")
+    parser.add_argument("--shape_latent_name", type=str, default="shape_enc_next_dc_f16c32_fp16_1024_view",
+                        help="Shape latent directory name under shape_latents/")
+    parser.add_argument("--pbr_latent_name", type=str, default="tex_enc_next_dc_f16c32_fp16_1024_view_fix",
+                        help="PBR latent directory name under pbr_latents/")
+    parser.add_argument("--shape_decoder", type=str, default="microsoft/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16",
+                        help="Pretrained shape decoder")
+    parser.add_argument("--pbr_decoder", type=str, default="microsoft/TRELLIS.2-4B/ckpts/tex_dec_next_dc_f16c32_fp16",
+                        help="Pretrained PBR/texture decoder")
+    parser.add_argument("--texture_size", type=int, default=4096, help="GLB texture resolution")
+    parser.add_argument("--decimation_target", type=int, default=1000000, help="GLB mesh decimation target")
+    parser.add_argument("--output_dir", type=str, default=None, help="Output directory (default: <root>/vis_pbr/<sha256>)")
+    args = parser.parse_args()
+
+    sha256 = args.sha256
+    root = args.root
+    view_idx = args.view_idx
+
+    # Paths
+    shape_latent_dir = os.path.join(root, "shape_latents", args.shape_latent_name, sha256)
+    pbr_latent_dir = os.path.join(root, "pbr_latents", args.pbr_latent_name, sha256)
+    shape_file = os.path.join(shape_latent_dir, f"view{view_idx:02d}.npz")
+    pbr_file = os.path.join(pbr_latent_dir, f"view{view_idx:02d}.npz")
+    renders_dir = os.path.join(root, "renders_cond", sha256)
+    output_dir = args.output_dir or os.path.join(root, "vis_pbr", sha256)
+
+    # Validate
+    assert os.path.exists(shape_file), f"Shape latent not found: {shape_file}"
+    assert os.path.exists(pbr_file), f"PBR latent not found: {pbr_file}"
+    print(f"[Input] Shape latent: {shape_file}")
+    print(f"[Input] PBR latent:   {pbr_file}")
+    if os.path.exists(renders_dir):
+        print(f"[Input] Renders:      {renders_dir}")
+
+    # 1. Load latents
+    print("[Step 1] Loading latents...")
+    shape_slat = load_latent(shape_file)
+    pbr_slat = load_latent(pbr_file)
+    print(f"  Shape: coords {shape_slat.coords.shape}, feats {shape_slat.feats.shape}")
+    print(f"  PBR:   coords {pbr_slat.coords.shape}, feats {pbr_slat.feats.shape}")
+
+    # 2. Load decoders
+    print(f"[Step 2] Loading decoders...")
+    shape_dec = models.from_pretrained(args.shape_decoder)
+    shape_dec.set_resolution(args.resolution)
+    shape_dec = shape_dec.cuda().eval()
+
+    pbr_dec = models.from_pretrained(args.pbr_decoder)
+    pbr_dec = pbr_dec.cuda().eval()
+
+    # 3. Decode shape → mesh + subs, then PBR → voxel
+    print("[Step 3] Decoding shape + PBR latents...")
+    with torch.no_grad():
+        meshes, subs = shape_dec(shape_slat, return_subs=True)
+        vox = pbr_dec(pbr_slat, guide_subs=subs) * 0.5 + 0.5
+
+    mesh = meshes[0]
+    mesh.fill_holes()
+    mesh_with_voxel = MeshWithVoxel(
+        mesh.vertices, mesh.faces,
+        origin=[-0.5, -0.5, -0.5],
+        voxel_size=1 / args.resolution,
+        coords=vox[0].coords[:, 1:],
+        attrs=vox[0].feats,
+        voxel_shape=torch.Size([*vox[0].shape, *vox[0].spatial_shape]),
+        layout=PBR_ATTR_LAYOUT,
+    )
+    print(f"  Mesh: vertices {mesh.vertices.shape}, faces {mesh.faces.shape}")
+    print(f"  Voxel: coords {vox[0].coords.shape}, feats {vox[0].feats.shape}")
+
+    # 4. Export GLB with PBR textures
+    print("[Step 4] Extracting textured GLB...")
+    os.makedirs(output_dir, exist_ok=True)
+    glb = o_voxel.postprocess.to_glb(
+        vertices=mesh_with_voxel.vertices,
+        faces=mesh_with_voxel.faces,
+        attr_volume=mesh_with_voxel.attrs,
+        coords=mesh_with_voxel.coords,
+        attr_layout=PBR_ATTR_LAYOUT,
+        grid_size=args.resolution,
+        aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
+        decimation_target=args.decimation_target,
+        texture_size=args.texture_size,
+        remesh=True, remesh_band=1, remesh_project=0, use_tqdm=True,
+    )
+    # Apply rotation (same as inference.py)
+    rot = np.array([
+        [-1,  0,  0,  0],
+        [ 0,  0, -1,  0],
+        [ 0, -1,  0,  0],
+        [ 0,  0,  0,  1],
+    ], dtype=np.float64)
+    glb.apply_transform(rot)
+    glb_path = os.path.join(output_dir, f"pbr_view{view_idx:02d}.glb")
+    glb.export(glb_path, extension_webp=True)
+    print(f"  GLB saved: {glb_path}")
+
+    # 5. Render PBR front view (proj-aligned, same as app.py)
+    print("[Step 5] Rendering PBR front view (proj-aligned)...")
+    transforms_file = os.path.join(renders_dir, "transforms.json")
+    shape_scale_file = os.path.join(shape_latent_dir, f"view{view_idx:02d}_scale.json")
+    envmaps = load_envmaps(device='cuda')
+    if os.path.exists(transforms_file) and os.path.exists(shape_scale_file) and envmaps:
+        with open(transforms_file) as f:
+            transforms = json.load(f)
+        with open(shape_scale_file) as f:
+            scale_info = json.load(f)
+        total_scale = scale_info['total_scale']
+        frame_info = transforms['frames'][view_idx]
+        camera_angle_x = frame_info['camera_angle_x']
+        distance = frame_info['radius']
+        near = max(0.01, distance - 2.0)
+        far = distance + 10.0
+        # Scale mesh by 1/total_scale to match blender normalized space
+        scaled_mesh = MeshWithVoxel(
+            mesh_with_voxel.vertices / total_scale,
+            mesh_with_voxel.faces,
+            origin=[x / total_scale for x in mesh_with_voxel.origin],
+            voxel_size=mesh_with_voxel.voxel_size / total_scale,
+            coords=mesh_with_voxel.coords,
+            attrs=mesh_with_voxel.attrs,
+            voxel_shape=mesh_with_voxel.voxel_shape,
+            layout=PBR_ATTR_LAYOUT,
+        )
+        print(f"  total_scale={total_scale:.4f}, distance={distance:.4f}, fov={camera_angle_x:.4f}")
+        renders = render_utils.render_proj_aligned_video(
+            scaled_mesh, camera_angle_x=camera_angle_x, distance=distance,
+            resolution=1024, num_frames=1, envmap=envmaps, near=near, far=far,
+        )
+        for key, frames in renders.items():
+            for i, frame in enumerate(frames):
+                img = Image.fromarray(frame)
+                img_path = os.path.join(output_dir, f"decoded_{key}_view{view_idx:02d}_{i:03d}.png")
+                img.save(img_path)
+            print(f"  Saved {len(frames)} {key} images")
+    else:
+        if not os.path.exists(transforms_file):
+            print("  No transforms.json found, skipping rendering.")
+        if not os.path.exists(shape_scale_file):
+            print("  No scale file found, skipping rendering.")
+        if not envmaps:
+            print("  No HDRI envmaps found, skipping PBR rendering.")
+
+    # Free GPU
+    del shape_dec, pbr_dec, shape_slat, pbr_slat, meshes, subs, vox
+    torch.cuda.empty_cache()
+
+    # 6. Copy condition renders
+    if os.path.exists(renders_dir):
+        print("[Step 6] Copying condition renders...")
+        for fname in sorted(os.listdir(renders_dir)):
+            src = os.path.join(renders_dir, fname)
+            dst = os.path.join(output_dir, fname)
+            shutil.copy2(src, dst)
+            print(f"  {fname}")
+    else:
+        print("[Step 6] No condition renders found, skipping.")
+
+    # 7. Copy scale info
+    for src_dir, prefix in [(shape_latent_dir, "shape"), (pbr_latent_dir, "pbr")]:
+        scale_file = os.path.join(src_dir, f"view{view_idx:02d}_scale.json")
+        if os.path.exists(scale_file):
+            shutil.copy2(scale_file, os.path.join(output_dir, f"{prefix}_view{view_idx:02d}_scale.json"))
+
+    print(f"\n[Done] All outputs in: {output_dir}")
+    print(f"  Files: {sorted(os.listdir(output_dir))}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/data_toolkit/visualize_shape_latent.py b/data_toolkit/visualize_shape_latent.py
new file mode 100644
index 0000000..4a39448
--- /dev/null
+++ b/data_toolkit/visualize_shape_latent.py
@@ -0,0 +1,160 @@
+"""
+Decode a view-aligned shape latent (.npz) to GLB mesh and render the front view.
+
+Usage:
+    python data_toolkit/visualize_shape_latent.py \
+        --root datasets/ObjaverseXL_sketchfab \
+        --sha256 <SHA256_HASH> \
+        --resolution 1024 \
+        --view_idx 0
+"""
+
+import os
+import sys
+import json
+import shutil
+import argparse
+import numpy as np
+import torch
+import trimesh
+from PIL import Image
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
+
+import pixal3d.models as models
+import pixal3d.modules.sparse as sp
+from pixal3d.utils import render_utils
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Decode shape latent to GLB and collect renders")
+    parser.add_argument("--root", type=str, required=True, help="Dataset root, e.g. /local-ssd/datasets/ObjaverseXL_sketchfab")
+    parser.add_argument("--sha256", type=str, required=True, help="SHA256 of the asset")
+    parser.add_argument("--resolution", type=int, default=1024, help="Decoder resolution (must match latent resolution)")
+    parser.add_argument("--view_idx", type=int, default=0, help="View index to decode")
+    parser.add_argument("--latent_name", type=str, default="shape_enc_next_dc_f16c32_fp16_1024_view",
+                        help="Latent directory name under shape_latents/")
+    parser.add_argument("--decoder", type=str, default="microsoft/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16",
+                        help="Pretrained shape decoder path (HuggingFace or local)")
+    parser.add_argument("--output_dir", type=str, default=None, help="Output directory (default: <root>/vis/<sha256>)")
+    args = parser.parse_args()
+
+    sha256 = args.sha256
+    root = args.root
+    view_idx = args.view_idx
+
+    # Paths
+    latent_dir = os.path.join(root, "shape_latents", args.latent_name, sha256)
+    latent_file = os.path.join(latent_dir, f"view{view_idx:02d}.npz")
+    scale_file = os.path.join(latent_dir, f"view{view_idx:02d}_scale.json")
+    renders_dir = os.path.join(root, "renders_cond", sha256)
+    output_dir = args.output_dir or os.path.join(root, "vis", sha256)
+
+    # Validate
+    assert os.path.exists(latent_file), f"Latent file not found: {latent_file}"
+    print(f"[Input] Latent: {latent_file}")
+    if os.path.exists(scale_file):
+        print(f"[Input] Scale:  {scale_file}")
+    if os.path.exists(renders_dir):
+        print(f"[Input] Renders: {renders_dir}")
+
+    # 1. Load latent
+    print("[Step 1] Loading shape latent...")
+    data = np.load(latent_file)
+    coords = torch.tensor(data['coords']).int()
+    feats = torch.tensor(data['feats']).float()
+    # Prepend batch dim (0) to coords
+    coords = torch.cat([torch.zeros_like(coords[:, :1]), coords], dim=1)
+    slat = sp.SparseTensor(feats.cuda(), coords.cuda())
+    print(f"  coords: {coords.shape}, feats: {feats.shape}")
+
+    # 2. Load decoder
+    print(f"[Step 2] Loading shape decoder: {args.decoder}")
+    decoder = models.from_pretrained(args.decoder)
+    decoder.set_resolution(args.resolution)
+    decoder = decoder.cuda().eval()
+
+    # 3. Decode
+    print("[Step 3] Decoding shape latent → mesh...")
+    with torch.no_grad():
+        meshes, subs = decoder(slat, return_subs=True)
+    mesh = meshes[0]
+    print(f"  vertices: {mesh.vertices.shape}, faces: {mesh.faces.shape}")
+
+    # 4. Convert to trimesh and export GLB
+    print("[Step 4] Exporting GLB...")
+    vertices = mesh.vertices.cpu().numpy()
+    faces = mesh.faces.cpu().numpy()
+
+    # Apply coordinate rotation (same as inference.py)
+    # Swap axes: x→-x, y→-z, z→-y
+    rot = np.array([
+        [-1,  0,  0],
+        [ 0,  0, -1],
+        [ 0, -1,  0],
+    ], dtype=np.float64)
+    vertices = vertices @ rot.T
+
+    tri_mesh = trimesh.Trimesh(vertices=vertices, faces=faces, process=False)
+
+    os.makedirs(output_dir, exist_ok=True)
+    glb_path = os.path.join(output_dir, f"shape_view{view_idx:02d}.glb")
+    tri_mesh.export(glb_path)
+    print(f"  GLB saved: {glb_path}")
+
+    # 5. Render front view (proj-aligned, same as app.py)
+    print("[Step 5] Rendering decoded mesh (proj-aligned front view)...")
+    transforms_file = os.path.join(renders_dir, "transforms.json")
+    if os.path.exists(transforms_file) and os.path.exists(scale_file):
+        with open(transforms_file) as f:
+            transforms = json.load(f)
+        with open(scale_file) as f:
+            scale_info = json.load(f)
+        total_scale = scale_info['total_scale']
+        frame_info = transforms['frames'][view_idx]
+        camera_angle_x = frame_info['camera_angle_x']
+        distance = frame_info['radius']
+        near = max(0.01, distance - 2.0)
+        far = distance + 10.0
+        # Scale mesh vertices by 1/total_scale to match blender normalized space
+        from pixal3d.representations import Mesh
+        scaled_mesh = Mesh(mesh.vertices / total_scale, mesh.faces)
+        print(f"  total_scale={total_scale:.4f}, distance={distance:.4f}, fov={camera_angle_x:.4f}")
+        renders = render_utils.render_proj_aligned_video(
+            scaled_mesh, camera_angle_x=camera_angle_x, distance=distance,
+            resolution=1024, num_frames=1, near=near, far=far,
+        )
+        for key, frames in renders.items():
+            for i, frame in enumerate(frames):
+                img = Image.fromarray(frame)
+                img_path = os.path.join(output_dir, f"decoded_{key}_view{view_idx:02d}_{i:03d}.png")
+                img.save(img_path)
+            print(f"  Saved {len(frames)} {key} images")
+    else:
+        print("  No transforms.json or scale file found, skipping rendering.")
+
+    # Free decoder GPU memory
+    del decoder, slat, meshes, subs
+    torch.cuda.empty_cache()
+
+    # 6. Copy renders
+    if os.path.exists(renders_dir):
+        print("[Step 6] Copying condition renders...")
+        for fname in sorted(os.listdir(renders_dir)):
+            src = os.path.join(renders_dir, fname)
+            dst = os.path.join(output_dir, fname)
+            shutil.copy2(src, dst)
+            print(f"  {fname}")
+    else:
+        print("[Step 6] No condition renders found, skipping.")
+
+    # 7. Copy scale info
+    if os.path.exists(scale_file):
+        shutil.copy2(scale_file, os.path.join(output_dir, f"view{view_idx:02d}_scale.json"))
+
+    print(f"\n[Done] All outputs in: {output_dir}")
+    print(f"  Files: {sorted(os.listdir(output_dir))}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/data_toolkit/voxelize_pbr.py b/data_toolkit/voxelize_pbr.py
new file mode 100644
index 0000000..23b75e2
--- /dev/null
+++ b/data_toolkit/voxelize_pbr.py
@@ -0,0 +1,167 @@
+import os
+import copy
+import sys
+import importlib
+import argparse
+import pandas as pd
+import pickle
+import numpy as np
+import torch
+from easydict import EasyDict as edict
+from functools import partial
+import o_voxel
+
+
+def _pbr_voxelize(file, metadatum, pbr_dump_root, root):
+    sha256 = metadatum['sha256']
+    try:
+        pack = {'sha256': sha256}
+        dump = None
+        for res in opt.resolution:
+            need_process = False
+
+            # check if already processed
+            if os.path.exists(os.path.join(root, f'pbr_voxels_{res}', f'{sha256}.vxz')):
+                try:
+                    info = o_voxel.io.read_vxz_info(os.path.join(root, f'pbr_voxels_{res}', f'{sha256}.vxz'))
+                    pack[f'pbr_voxelized_{res}'] = True
+                    pack[f'num_pbr_voxels_{res}'] = info['num_voxel']
+                except Exception as e:
+                    print(f'Error reading {sha256}.vxz: {e}')
+                    need_process = True
+            else:
+                need_process = True
+
+            # process if necessary
+            if need_process:
+                if dump == None:
+                    with open(os.path.join(pbr_dump_root, 'pbr_dumps', f'{sha256}.pickle'), 'rb') as f:
+                        dump = pickle.load(f)
+                    # Fix dump alpha map
+                    for mat in dump['materials']:
+                        if mat['alphaTexture'] is not None and mat['alphaMode'] == 'OPAQUE':
+                            mat['alphaMode'] = 'BLEND'
+                    dump['materials'].append({
+                        "baseColorFactor": [0.8, 0.8, 0.8],
+                        "alphaFactor": 1.0,
+                        "metallicFactor": 0.0,
+                        "roughnessFactor": 0.5,
+                        "alphaMode": "OPAQUE",
+                        "alphaCutoff": 0.5,
+                        "baseColorTexture": None,
+                        "alphaTexture": None,
+                        "metallicTexture": None,
+                        "roughnessTexture": None,
+                    })      # append default material
+                    dump['objects'] = [
+                        obj for obj in dump['objects']
+                        if obj['vertices'].size != 0 and obj['faces'].size != 0
+                    ]
+                    vertices = torch.from_numpy(np.concatenate([obj['vertices'] for obj in dump['objects']], axis=0)).float()
+                    vertices_min = vertices.min(dim=0)[0]
+                    vertices_max = vertices.max(dim=0)[0]
+                    center = (vertices_min + vertices_max) / 2
+                    scale = 0.99999 / (vertices_max - vertices_min).max()
+                    for obj in dump['objects']:
+                        obj['vertices'] = (torch.from_numpy(obj['vertices']).float() - center) * scale
+                        obj['vertices'] = obj['vertices'].numpy()
+                        obj['mat_ids'][obj['mat_ids'] == -1] = len(dump['materials']) - 1
+                        assert np.all(obj['mat_ids'] >= 0), 'invalid mat_ids'
+                        assert np.all(obj['vertices'] >= -0.5) and np.all(obj['vertices'] <= 0.5), 'vertices out of range'
+
+                coord, attr = o_voxel.convert.blender_dump_to_volumetric_attr(dump, grid_size=res, aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
+                                                                              mip_level_offset=0, verbose=False, timing=False)
+                del attr['normal']
+                del attr['emissive']
+                o_voxel.io.write_vxz(os.path.join(root, f'pbr_voxels_{res}', f'{sha256}.vxz'), coord, attr)
+                pack[f'pbr_voxelized_{res}'] = True
+                pack[f'num_pbr_voxels_{res}'] = len(coord)
+
+        return pack
+    except Exception as e:
+        print(f'Error voxelizing {sha256}: {e}')
+        return {'sha256': sha256, 'error': str(e)}
+
+
+if __name__ == '__main__':
+    dataset_utils = importlib.import_module(f'datasets.{sys.argv[1]}')
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--root', type=str, required=True,
+                        help='Directory to save the metadata')
+    parser.add_argument('--pbr_dump_root', type=str, default=None,
+                        help='Directory to load mesh dumps')
+    parser.add_argument('--pbr_voxel_root', type=str, default=None,
+                        help='Directory to save voxelized pbr attributes')
+    parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
+                        help='Filter objects with aesthetic score lower than this value')
+    parser.add_argument('--instances', type=str, default=None,
+                        help='Instances to process')
+    dataset_utils.add_args(parser)
+    parser.add_argument('--resolution', type=str, default=1024)
+    parser.add_argument('--rank', type=int, default=0)
+    parser.add_argument('--world_size', type=int, default=1)
+    parser.add_argument('--max_workers', type=int, default=0)
+    opt = parser.parse_args(sys.argv[2:])
+    opt = edict(vars(opt))
+    opt.resolution = sorted([int(x) for x in opt.resolution.split(',')], reverse=True)
+    opt.pbr_dump_root = opt.pbr_dump_root or opt.root
+    opt.pbr_voxel_root = opt.pbr_voxel_root or opt.root
+
+    for res in opt.resolution:
+        os.makedirs(os.path.join(opt.pbr_voxel_root, f'pbr_voxels_{res}', 'new_records'), exist_ok=True)
+
+    # get file list
+    if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
+        raise ValueError('metadata.csv not found')
+    metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
+    if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256'))
+    if os.path.exists(os.path.join(opt.pbr_dump_root, 'pbr_dumps', 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.pbr_dump_root, 'pbr_dumps', 'metadata.csv')).set_index('sha256'))
+    for res in opt.resolution:
+        if os.path.exists(os.path.join(opt.pbr_voxel_root, f'pbr_voxels_{res}', 'metadata.csv')):
+            pbr_voxel_metadata = pd.read_csv(os.path.join(opt.pbr_voxel_root, f'pbr_voxels_{res}','metadata.csv')).set_index('sha256')
+            pbr_voxel_metadata = pbr_voxel_metadata.rename(columns={'pbr_voxelized': f'pbr_voxelized_{res}', 'num_pbr_voxels': f'num_pbr_voxels_{res}'})
+            metadata = metadata.combine_first(pbr_voxel_metadata)
+    metadata = metadata.reset_index()
+    if opt.instances is None:
+        if opt.filter_low_aesthetic_score is not None:
+            metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
+        metadata = metadata[metadata['pbr_dumped'] == True]
+        mask = np.zeros(len(metadata), dtype=bool)
+        for res in opt.resolution:
+            if f'pbr_voxelized_{res}' in metadata.columns:
+                mask |= metadata[f'pbr_voxelized_{res}'] != True
+            else:
+                mask[:] = True
+                break
+        metadata = metadata[mask]
+    else:
+        if os.path.exists(opt.instances):
+            with open(opt.instances, 'r') as f:
+                instances = f.read().splitlines()
+        else:
+            instances = opt.instances.split(',')
+        metadata = metadata[metadata['sha256'].isin(instances)]
+
+    start = len(metadata) * opt.rank // opt.world_size
+    end = len(metadata) * (opt.rank + 1) // opt.world_size
+    metadata = metadata[start:end]
+    
+    print(f'Processing {len(metadata)} objects...')
+
+    # process objects
+    func = partial(_pbr_voxelize, pbr_dump_root=opt.pbr_dump_root, root=opt.pbr_voxel_root)
+    pbr_voxelized = dataset_utils.foreach_instance(metadata, None, func, max_workers=opt.max_workers, no_file=True, desc='Voxelizing')
+    if 'error' in pbr_voxelized.columns:
+        errors = pbr_voxelized[pbr_voxelized['error'].notna()]
+        with open('errors.txt', 'w') as f:
+            f.write('\n'.join(errors['sha256'].tolist()))
+    for res in opt.resolution:
+        if f'pbr_voxelized_{res}' in pbr_voxelized.columns:
+            pbr_voxel_metadata = pbr_voxelized[pbr_voxelized[f'pbr_voxelized_{res}'] == True]
+            if len(pbr_voxel_metadata) > 0:
+                pbr_voxel_metadata = pbr_voxel_metadata[['sha256', f'pbr_voxelized_{res}', f'num_pbr_voxels_{res}']]
+                pbr_voxel_metadata = pbr_voxel_metadata.rename(columns={f'pbr_voxelized_{res}': 'pbr_voxelized', f'num_pbr_voxels_{res}': 'num_pbr_voxels'})
+                pbr_voxel_metadata.to_csv(os.path.join(opt.pbr_voxel_root, f'pbr_voxels_{res}', 'new_records', f'part_{opt.rank}.csv'), index=False)
diff --git a/data_toolkit/voxelize_pbr_view.py b/data_toolkit/voxelize_pbr_view.py
new file mode 100644
index 0000000..11fea8f
--- /dev/null
+++ b/data_toolkit/voxelize_pbr_view.py
@@ -0,0 +1,596 @@
+"""
+voxelize_pbr_view.py - Multi-view transform PBR voxelization
+Extends voxelize_pbr.py with scale and mesh rotation logic
+Based on dual_grid_view.py and test_ovoxel_pbr_transform.py implementation
+"""
+import os
+import copy
+import sys
+import importlib
+import argparse
+import json
+import math
+import pandas as pd
+import pickle
+import numpy as np
+import torch
+from easydict import EasyDict as edict
+from functools import partial
+import o_voxel
+from utils import get_new_camera_matrix, sphere_normalize_torch
+
+
+# ==================== PBR-specific transform functions ====================
+
+def transform_vertices(vertices, frame):
+    """
+    Apply multi-view transform to vertices based on camera transform matrix.
+
+    Args:
+        vertices: torch.Tensor, shape [N, 3], vertex coordinates
+        frame: dict containing transform_matrix
+
+    Returns:
+        transformed_vertices: torch.Tensor, shape [N, 3]
+    """
+    device = vertices.device
+    c2w_orig = torch.tensor(frame['transform_matrix'], dtype=torch.float32, device=device)
+
+    # Old and new camera matrices
+    radius = c2w_orig[:3, 3].norm().item()
+    c2w_new = get_new_camera_matrix(radius=radius, yaw=-90/180.0*math.pi, pitch=0.0,
+                                dtype=torch.float32, device=device)
+    w2c_orig = torch.inverse(c2w_orig)
+
+    # Initial and final axis alignment matrices
+    R_init = torch.tensor([
+        [1.0, 0.0,  0.0, 0.0],
+        [0.0, 0.0, -1.0, 0.0],
+        [0.0, 1.0,  0.0, 0.0],
+        [0.0, 0.0,  0.0, 1.0]
+    ], dtype=torch.float32, device=device)
+
+    R_back = torch.tensor([
+        [1.0,  0.0, 0.0, 0.0],
+        [0.0,  0.0, 1.0, 0.0],
+        [0.0, -1.0, 0.0, 0.0],
+        [0.0,  0.0, 0.0, 1.0]
+    ], dtype=torch.float32, device=device)
+
+    R_ply = torch.tensor([
+        [1.0,  0.0, 0.0, 0.0],
+        [0.0,  0.0, 1.0, 0.0],
+        [0.0, -1.0, 0.0, 0.0],
+        [0.0,  0.0, 0.0, 1.0]
+    ], dtype=torch.float32, device=device)
+
+    T_cam = c2w_new @ w2c_orig @ R_ply
+    T_final = R_back @ T_cam @ R_init
+
+    # Apply transform
+    vertices = vertices.reshape(-1, 3)
+    verts_h = torch.cat([vertices, torch.ones((vertices.shape[0], 1), dtype=torch.float32, device=device)], dim=1)
+    verts_trans = (T_final @ verts_h.T).T[:, :3]
+
+    return verts_trans
+
+
+def transform_normals(normals, frame):
+    """
+    Apply multi-view transform to normals (rotation only).
+    Consistent with test_ovoxel_pbr_transform.py implementation.
+
+    Args:
+        normals: torch.Tensor or np.ndarray, shape [N, 3] or [N, 3, 3]
+        frame: dict containing transform_matrix
+
+    Returns:
+        transformed_normals: np.ndarray (always returns numpy for dump compatibility)
+    """
+    is_numpy = isinstance(normals, np.ndarray)
+    if is_numpy:
+        normals = torch.from_numpy(normals).float()
+
+    device = normals.device
+    original_shape = normals.shape
+
+    # Flatten to [N, 3] for processing
+    if len(original_shape) == 3:
+        normals_flat = normals.reshape(-1, 3)
+    else:
+        normals_flat = normals
+
+    c2w_orig = torch.tensor(frame['transform_matrix'], dtype=torch.float32, device=device)
+
+    # Old and new camera matrices
+    radius = c2w_orig[:3, 3].norm().item()
+    c2w_new = get_new_camera_matrix(radius=radius, yaw=-90/180.0*math.pi, pitch=0.0,
+                                dtype=torch.float32, device=device)
+    w2c_orig = torch.inverse(c2w_orig)
+
+    # Axis alignment matrices (rotation part only, 3x3)
+    R_init = torch.tensor([
+        [1.0, 0.0,  0.0],
+        [0.0, 0.0, -1.0],
+        [0.0, 1.0,  0.0]
+    ], dtype=torch.float32, device=device)
+
+    R_back = torch.tensor([
+        [1.0,  0.0, 0.0],
+        [0.0,  0.0, 1.0],
+        [0.0, -1.0, 0.0]
+    ], dtype=torch.float32, device=device)
+
+    R_ply = torch.tensor([
+        [1.0,  0.0, 0.0],
+        [0.0,  0.0, 1.0],
+        [0.0, -1.0, 0.0]
+    ], dtype=torch.float32, device=device)
+
+    # Use rotation part only
+    T_cam_rot = c2w_new[:3, :3] @ w2c_orig[:3, :3] @ R_ply
+    T_final_rot = R_back @ T_cam_rot @ R_init
+
+    # Apply rotation transform
+    normals_trans = torch.matmul(normals_flat, T_final_rot.T)
+
+    # Re-normalize
+    normals_trans = torch.nn.functional.normalize(normals_trans, dim=-1)
+
+    # Restore original shape
+    if len(original_shape) == 3:
+        normals_trans = normals_trans.reshape(original_shape)
+
+    # Always return numpy array for dump compatibility
+    return normals_trans.numpy()
+
+
+def prepare_pbr_dump(dump):
+    """
+    Prepare PBR dump data for processing.
+    Consistent with voxelize_pbr.py preprocessing.
+
+    Args:
+        dump: raw PBR dump data
+
+    Returns:
+        processed dump data
+    """
+    dump = copy.deepcopy(dump)
+
+    # Fix dump alpha map
+    for mat in dump['materials']:
+        if mat['alphaTexture'] is not None and mat['alphaMode'] == 'OPAQUE':
+            mat['alphaMode'] = 'BLEND'
+
+    # Append default material
+    dump['materials'].append({
+        "baseColorFactor": [0.8, 0.8, 0.8],
+        "alphaFactor": 1.0,
+        "metallicFactor": 0.0,
+        "roughnessFactor": 0.5,
+        "alphaMode": "OPAQUE",
+        "alphaCutoff": 0.5,
+        "baseColorTexture": None,
+        "alphaTexture": None,
+        "metallicTexture": None,
+        "roughnessTexture": None,
+    })
+
+    # Filter out empty objects
+    dump['objects'] = [
+        obj for obj in dump['objects']
+        if obj['vertices'].size != 0 and obj['faces'].size != 0
+    ]
+
+    return dump
+
+
+def transform_pbr_dump(dump, frame):
+    """
+    Apply multi-view transform to entire PBR dump data.
+
+    Processing flow (based on test_ovoxel_pbr_transform.py):
+    1. Box normalize all vertices (scale only, no center shift)
+    2. Sphere normalize
+    3. Apply multi-view transform
+    4. Normalize back to [-0.5, 0.5]^3
+
+    Note: All object vertices are processed together (not per-object) for consistency.
+
+    Args:
+        dump: PBR dump data (already preprocessed via prepare_pbr_dump)
+        frame: camera frame info
+
+    Returns:
+        transformed_dump: transformed dump data
+        total_scale: total scale from original mesh to final mesh
+    """
+    transformed_dump = copy.deepcopy(dump)
+
+    # 1. Collect all vertices
+    all_vertices_list = []
+    vertex_counts = []
+    for obj in transformed_dump['objects']:
+        all_vertices_list.append(obj['vertices'])
+        vertex_counts.append(len(obj['vertices']))
+
+    if len(all_vertices_list) == 0:
+        return transformed_dump, 1.0
+
+    all_vertices = np.concatenate(all_vertices_list, axis=0)
+
+    # 2. Box normalize (scale only, no center shift, consistent with original rendering)
+    vertices_min = all_vertices.min(axis=0)
+    vertices_max = all_vertices.max(axis=0)
+    box_scale_init = 0.99999 / (vertices_max - vertices_min).max()
+    all_vertices_box_normalized = all_vertices * box_scale_init
+
+    all_vertices_tensor = torch.from_numpy(all_vertices_box_normalized).float()
+
+    # 3. Sphere normalize all vertices together
+    all_vertices_sphere, sphere_center, sphere_radius = sphere_normalize_torch(all_vertices_tensor)
+
+    # 4. Multi-view transform
+    all_transformed = transform_vertices(all_vertices_sphere, frame)
+
+    # 5. Normalize back to [-0.5, 0.5]^3 (all vertices together)
+    abs_max = all_transformed.abs().max().item()
+    box_scale_final = 0.49999 / abs_max
+    all_transformed_normalized = all_transformed * box_scale_final
+
+    # Compute total scale (from original mesh to final normalized mesh)
+    total_scale = box_scale_init * box_scale_final / sphere_radius.item()
+
+    # 6. Split back to individual objects
+    start_idx = 0
+    for i, obj in enumerate(transformed_dump['objects']):
+        end_idx = start_idx + vertex_counts[i]
+        obj['vertices'] = all_transformed_normalized[start_idx:end_idx].numpy()
+        start_idx = end_idx
+
+        # Transform normals
+        if obj['normals'] is not None and obj['normals'].size > 0:
+            obj['normals'] = transform_normals(obj['normals'], frame)
+
+        # Fix mat_ids (replace -1 with default material index)
+        obj['mat_ids'][obj['mat_ids'] == -1] = len(transformed_dump['materials']) - 1
+
+        # Validate range
+        assert np.all(obj['mat_ids'] >= 0), 'invalid mat_ids'
+        assert np.all(obj['vertices'] >= -0.5) and np.all(obj['vertices'] <= 0.5), 'vertices out of range'
+
+    return transformed_dump, total_scale
+
+
+def _pbr_voxelize_view(file, sha256, pbr_dump_root, transform_root, root, view_indices=None):
+    """
+    Process multi-view PBR voxelization for a single sha256.
+
+    Args:
+        file: local_path from metadata
+        sha256: sha256 string
+        pbr_dump_root: directory containing PBR dump files
+        transform_root: directory containing transform json files
+        root: output directory for PBR voxels
+        view_indices: list of view indices to process, None for all views
+    """
+    try:
+        pack = {'sha256': sha256}
+        dump = None
+
+        # Load transforms
+        transform_path = os.path.join(transform_root, sha256, 'transforms.json')
+        if not os.path.exists(transform_path):
+            print(f'Transform file not found for {sha256}, skipping')
+            return {'sha256': sha256, 'error': 'Transform file not found'}
+
+        with open(transform_path, 'r') as f:
+            transforms_json = json.load(f)
+        transform_mats = transforms_json['frames']
+
+        # Determine views to process
+        if view_indices is None:
+            view_indices = list(range(len(transform_mats)))
+        else:
+            view_indices = [i for i in view_indices if i < len(transform_mats)]
+
+        # Track processed and skipped counts
+        processed_count = 0
+        skipped_count = 0
+
+        for view_idx in view_indices:
+            for res in opt.resolution:
+                need_process = False
+
+                # Check if already processed
+                # Path structure: pbr_voxels_view_fix_{res}/{sha256}/view{idx:02d}.vxz
+                sha256_dir = os.path.join(root, f'pbr_voxels_view_fix_{res}', sha256)
+                vxz_path = os.path.join(sha256_dir, f'view{view_idx:02d}.vxz')
+                if os.path.exists(vxz_path):
+                    try:
+                        info = o_voxel.io.read_vxz_info(vxz_path)
+                        pack[f'pbr_voxelized_view_fix{view_idx:02d}_{res}'] = True
+                        pack[f'num_pbr_voxels_view_fix{view_idx:02d}_{res}'] = info['num_voxel']
+                        skipped_count += 1
+                    except Exception as e:
+                        print(f'Error reading {sha256}/view{view_idx:02d}.vxz: {e}, will reprocess')
+                        need_process = True
+                else:
+                    need_process = True
+
+                # Process PBR dump
+                if need_process:
+                    # Lazy load dump
+                    if dump is None:
+                        pbr_dump_file = os.path.join(pbr_dump_root, 'pbr_dumps', f'{sha256}.pickle')
+                        if not os.path.exists(pbr_dump_file):
+                            print(f'PBR dump not found for {sha256}, skipping')
+                            return {'sha256': sha256, 'error': 'PBR dump not found'}
+
+                        with open(pbr_dump_file, 'rb') as f:
+                            dump = pickle.load(f)
+
+                        # Prepare dump data
+                        dump = prepare_pbr_dump(dump)
+
+                        if len(dump['objects']) == 0:
+                            print(f'No valid objects in PBR dump for {sha256}, skipping')
+                            return {'sha256': sha256, 'error': 'No valid objects in PBR dump'}
+
+                    # Get transform for current view
+                    frame = transform_mats[view_idx]
+
+                    # Multi-view transform (deep copy from original dump each time)
+                    transformed_dump, total_scale = transform_pbr_dump(dump, frame)
+
+                    # PBR voxelization
+                    coord, attr = o_voxel.convert.blender_dump_to_volumetric_attr(
+                        transformed_dump,
+                        grid_size=res,
+                        aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
+                        mip_level_offset=0,
+                        verbose=False,
+                        timing=False
+                    )
+
+                    # Remove normal and emissive (consistent with voxelize_pbr.py)
+                    del attr['normal']
+                    del attr['emissive']
+
+                    # Save .vxz file
+                    os.makedirs(sha256_dir, exist_ok=True)
+                    o_voxel.io.write_vxz(vxz_path, coord, attr)
+
+                    # Save scale info
+                    scale_path = os.path.join(sha256_dir, f'view{view_idx:02d}_scale.json')
+                    scale_info = {
+                        'sha256': sha256,
+                        'view_idx': view_idx,
+                        'total_scale': float(total_scale),
+                    }
+                    with open(scale_path, 'w') as f:
+                        json.dump(scale_info, f, indent=2)
+
+                    pack[f'pbr_voxelized_view_fix{view_idx:02d}_{res}'] = True
+                    pack[f'num_pbr_voxels_view_fix{view_idx:02d}_{res}'] = len(coord)
+                    pack[f'pbr_voxel_scale_view_fix{view_idx:02d}_{res}'] = float(total_scale)
+                    processed_count += 1
+
+        # Record processing stats
+        pack['_processed_count'] = processed_count
+        pack['_skipped_count'] = skipped_count
+
+        return pack
+
+    except Exception as e:
+        print(f'Error processing {sha256}: {e}')
+        import traceback
+        traceback.print_exc()
+        return {'sha256': sha256, 'error': str(e)}
+
+
+if __name__ == '__main__':
+    dataset_utils = importlib.import_module(f'datasets.{sys.argv[1]}')
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--root', type=str, required=True,
+                        help='Directory to save the metadata')
+    parser.add_argument('--pbr_dump_root', type=str, default=None,
+                        help='Directory to load PBR dumps')
+    parser.add_argument('--transform_root', type=str, default=None,
+                        help='Directory to load transform json files (renders_cond)')
+    parser.add_argument('--pbr_voxel_root', type=str, default=None,
+                        help='Directory to save voxelized PBR attributes')
+    parser.add_argument('--filter_low_aesthetic_score', type=float, default=None,
+                        help='Filter objects with aesthetic score lower than this value')
+    parser.add_argument('--instances', type=str, default=None,
+                        help='Instances to process')
+    parser.add_argument('--view_indices', type=str, default=None,
+                        help='View indices to process, e.g., "0,1,2" or "0-5". None for all views')
+    parser.add_argument('--skip_list', type=str, default=None,
+                        help='Path to a file containing sha256 hashes to skip (one per line). '
+                             'Supports format: "sha256" or "dataset/sha256"')
+    parser.add_argument('--clean_pbr_dir', type=str, default=None,
+                        help='Path to clean_pbr directory. Will auto-load {dataset}_clean_output.txt as ok-list, '
+                             'only sha256 in ok-list will be processed')
+    parser.add_argument('--clean_pbr_name', type=str, default=None,
+                        help='Dataset name prefix for clean_pbr file (e.g., ObjaverseXL_github). '
+                             'Defaults to sys.argv[1] if not specified')
+    dataset_utils.add_args(parser)
+    parser.add_argument('--resolution', type=str, default='1024')
+    parser.add_argument('--rank', type=int, default=0)
+    parser.add_argument('--world_size', type=int, default=1)
+    parser.add_argument('--max_workers', type=int, default=0)
+    opt = parser.parse_args(sys.argv[2:])
+    opt = edict(vars(opt))
+    opt.resolution = sorted([int(x) for x in opt.resolution.split(',')], reverse=True)
+    opt.pbr_dump_root = opt.pbr_dump_root or opt.root
+    opt.transform_root = opt.transform_root or os.path.join(opt.root, 'renders_cond')
+    opt.pbr_voxel_root = opt.pbr_voxel_root or opt.root
+
+    # Parse view_indices
+    view_indices = None
+    if opt.view_indices is not None:
+        view_indices = []
+        for part in opt.view_indices.split(','):
+            if '-' in part:
+                start, end = map(int, part.split('-'))
+                view_indices.extend(range(start, end + 1))
+            else:
+                view_indices.append(int(part))
+        view_indices = list(set(view_indices))  # Deduplicate
+        view_indices.sort()
+
+    # Load skip list (sha256 hashes to skip)
+    skip_set = set()
+    if opt.skip_list is not None and os.path.exists(opt.skip_list):
+        with open(opt.skip_list, 'r') as f:
+            for line in f:
+                line = line.strip()
+                if line and not line.startswith('#'):
+                    # Support "dataset/sha256" and plain "sha256" format, extract pure sha256
+                    skip_set.add(line.split('/')[-1])
+        print(f'Loaded {len(skip_set)} items from skip_list: {opt.skip_list}')
+
+    # Load clean_pbr ok-list (only process approved sha256)
+    ok_set = None
+    if opt.clean_pbr_dir is not None:
+        dataset_name = opt.clean_pbr_name or sys.argv[1]
+        clean_file = os.path.join(opt.clean_pbr_dir, f'{dataset_name}_clean_output.txt')
+        if os.path.exists(clean_file):
+            ok_set = set()
+            with open(clean_file, 'r') as f:
+                for line in f:
+                    line = line.strip()
+                    if line and not line.startswith('#'):
+                        ok_set.add(line.split('/')[-1])
+            print(f'Loaded {len(ok_set)} ok items from clean_pbr: {clean_file}')
+        else:
+            print(f'Warning: clean_pbr file not found: {clean_file}, proceeding without ok-list filter')
+
+    for res in opt.resolution:
+        os.makedirs(os.path.join(opt.pbr_voxel_root, f'pbr_voxels_view_fix_{res}', 'new_records'), exist_ok=True)
+
+    # Get file list
+    if not os.path.exists(os.path.join(opt.root, 'metadata.csv')):
+        raise ValueError('metadata.csv not found')
+    metadata = pd.read_csv(os.path.join(opt.root, 'metadata.csv')).set_index('sha256')
+    if os.path.exists(os.path.join(opt.root, 'aesthetic_scores', 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.root, 'aesthetic_scores','metadata.csv')).set_index('sha256'))
+    if os.path.exists(os.path.join(opt.pbr_dump_root, 'pbr_dumps', 'metadata.csv')):
+        metadata = metadata.combine_first(pd.read_csv(os.path.join(opt.pbr_dump_root, 'pbr_dumps', 'metadata.csv')).set_index('sha256'))
+
+    # Check already processed pbr_voxels_view_fix
+    for res in opt.resolution:
+        if os.path.exists(os.path.join(opt.pbr_voxel_root, f'pbr_voxels_view_fix_{res}', 'metadata.csv')):
+            pbr_voxel_metadata = pd.read_csv(os.path.join(opt.pbr_voxel_root, f'pbr_voxels_view_fix_{res}', 'metadata.csv')).set_index('sha256')
+            metadata = metadata.combine_first(pbr_voxel_metadata)
+
+    metadata = metadata.reset_index()
+
+    if opt.instances is None:
+        if opt.filter_low_aesthetic_score is not None:
+            metadata = metadata[metadata['aesthetic_score'] >= opt.filter_low_aesthetic_score]
+        metadata = metadata[metadata['pbr_dumped'] == True]
+
+        # Filter out objects with all views already processed
+        if view_indices is not None:
+            for res in opt.resolution:
+                # Check if each specified view is already processed
+                all_views_done_col = f'_all_views_done_{res}'
+                metadata[all_views_done_col] = True
+                for view_idx in view_indices:
+                    col_name = f'pbr_voxelized_view_fix{view_idx:02d}_{res}'
+                    if col_name in metadata.columns:
+                        metadata[all_views_done_col] = metadata[all_views_done_col] & (metadata[col_name] == True)
+                    else:
+                        metadata[all_views_done_col] = False
+                        break
+
+            # Keep objects with at least one incomplete resolution
+            any_incomplete = None
+            for res in opt.resolution:
+                all_views_done_col = f'_all_views_done_{res}'
+                if all_views_done_col in metadata.columns:
+                    if any_incomplete is None:
+                        any_incomplete = ~metadata[all_views_done_col]
+                    else:
+                        any_incomplete = any_incomplete | ~metadata[all_views_done_col]
+
+            if any_incomplete is not None:
+                before_filter = len(metadata)
+                metadata = metadata[any_incomplete]
+                print(f'Filtered out {before_filter - len(metadata)} already completed objects')
+    else:
+        if os.path.exists(opt.instances):
+            with open(opt.instances, 'r') as f:
+                instances = f.read().splitlines()
+        else:
+            instances = opt.instances.split(',')
+        metadata = metadata[metadata['sha256'].isin(instances)]
+
+    # Apply skip_list filter (exclude specified sha256)
+    if skip_set:
+        before_skip = len(metadata)
+        metadata = metadata[~metadata['sha256'].isin(skip_set)]
+        print(f'Skip list: filtered out {before_skip - len(metadata)} objects, {len(metadata)} remaining')
+
+    # Apply clean_pbr ok-list filter (only keep approved sha256)
+    if ok_set is not None:
+        before_ok = len(metadata)
+        metadata = metadata[metadata['sha256'].isin(ok_set)]
+        print(f'Ok list: kept {len(metadata)} objects out of {before_ok} (filtered {before_ok - len(metadata)})')
+
+    metadata = metadata.sample(frac=1, random_state=444).reset_index(drop=True)
+    start = len(metadata) * opt.rank // opt.world_size
+    end = len(metadata) * (opt.rank + 1) // opt.world_size
+    metadata = metadata[start:end]
+
+    print(f'Processing {len(metadata)} objects...')
+    if view_indices:
+        print(f'View indices to process: {view_indices}')
+    else:
+        print('Processing all available views')
+
+    # Process objects
+    func = partial(_pbr_voxelize_view,
+                   pbr_dump_root=opt.pbr_dump_root,
+                   transform_root=opt.transform_root,
+                   root=opt.pbr_voxel_root,
+                   view_indices=view_indices)
+    pbr_voxelized = dataset_utils.foreach_instance(metadata, opt.root, func, max_workers=opt.max_workers, desc='Voxelizing PBR views')
+
+    # Processing summary
+    total_processed = pbr_voxelized['_processed_count'].sum() if '_processed_count' in pbr_voxelized.columns else 0
+    total_skipped = pbr_voxelized['_skipped_count'].sum() if '_skipped_count' in pbr_voxelized.columns else 0
+    print(f'\n========== Processing Summary ==========')
+    print(f'Total processed (new): {int(total_processed)}')
+    print(f'Total skipped (existing): {int(total_skipped)}')
+    print(f'Total items: {int(total_processed + total_skipped)}')
+    print(f'=========================================\n')
+
+    if 'error' in pbr_voxelized.columns:
+        errors = pbr_voxelized[pbr_voxelized['error'].notna()]
+        if len(errors) > 0:
+            with open('errors_pbr_view.txt', 'w') as f:
+                f.write('\n'.join(errors['sha256'].tolist()))
+            print(f'Errors written to errors_pbr_view.txt ({len(errors)} errors)')
+
+    # Save metadata
+    for res in opt.resolution:
+        # Collect all view-related columns
+        view_cols = [col for col in pbr_voxelized.columns if f'pbr_voxelized_view_fix' in col and f'_{res}' in col]
+        if view_cols:
+            # Save metadata for each view
+            pbr_voxel_metadata = pbr_voxelized[pbr_voxelized[view_cols].any(axis=1)]
+            if len(pbr_voxel_metadata) > 0:
+                # Save simplified metadata
+                cols_to_save = ['sha256'] + [col for col in pbr_voxelized.columns if f'_{res}' in col]
+                cols_to_save = [col for col in cols_to_save if col in pbr_voxelized.columns]
+                pbr_voxel_metadata[cols_to_save].to_csv(
+                    os.path.join(opt.pbr_voxel_root, f'pbr_voxels_view_fix_{res}', 'new_records', f'part_{opt.rank}.csv'),
+                    index=False
+                )
+
+    print('Done!')
diff --git a/index.html b/index.html
new file mode 100644
index 0000000..99a9933
--- /dev/null
+++ b/index.html
@@ -0,0 +1,1493 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Pixal3D | AI Image-to-3D</title>
+    
+    <!-- Fonts & Icons -->
+    <link rel="preconnect" href="https://fonts.googleapis.com">
+    <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
+    <link href="https://fonts.googleapis.com/css2?family=Plus+Jakarta+Sans:wght@300;400;500;600;700;800&family=Outfit:wght@400;500;600;700;800&display=swap" rel="stylesheet">
+    <script src="https://unpkg.com/lucide@latest"></script>
+    <script type="module" src="https://ajax.googleapis.com/ajax/libs/model-viewer/4.0.0/model-viewer.min.js"></script>
+    
+    <style>
+        :root {
+            --primary: #818cf8;
+            --primary-dark: #6366f1;
+            --accent: #10b981;
+            --bg: #0b0f1a;
+            --surface: #161c2d;
+            --surface-light: #222b3e;
+            --border: rgba(255, 255, 255, 0.08);
+            --text: #f1f5f9;
+            --text-dim: #94a3b8;
+            --glass: rgba(255, 255, 255, 0.03);
+            --radius-lg: 24px;
+            --radius-md: 16px;
+            --radius-sm: 8px;
+        }
+
+        * {
+            margin: 0;
+            padding: 0;
+            box-sizing: border-box;
+        }
+
+        body {
+            font-family: 'Plus Jakarta Sans', sans-serif;
+            background: var(--bg);
+            color: var(--text);
+            min-height: 100vh;
+            display: flex;
+            flex-direction: column;
+            overflow-x: hidden;
+            background: 
+                radial-gradient(circle at 0% 0%, rgba(99, 102, 241, 0.15) 0%, transparent 40%),
+                radial-gradient(circle at 100% 100%, rgba(16, 185, 129, 0.1) 0%, transparent 40%);
+        }
+
+        /* Top Navigation / Steps */
+        .app-shell {
+            display: flex;
+            height: 100vh;
+            width: 100vw;
+        }
+
+        .sidebar {
+            width: 380px;
+            background: var(--surface);
+            border-right: 1px solid var(--border);
+            display: flex;
+            flex-direction: column;
+            padding: 1.5rem;
+            overflow-y: auto;
+            z-index: 10;
+        }
+
+        .main-content {
+            flex: 1;
+            min-width: 0;
+            display: flex;
+            flex-direction: column;
+            position: relative;
+            background: rgba(0,0,0,0.2);
+        }
+
+        header {
+            padding: 1rem 2rem;
+            display: flex;
+            align-items: center;
+            justify-content: space-between;
+            border-bottom: 1px solid var(--border);
+            background: rgba(11, 15, 26, 0.8);
+            backdrop-filter: blur(10px);
+        }
+
+        .logo {
+            display: flex;
+            align-items: center;
+            gap: 0.75rem;
+            font-family: 'Outfit', sans-serif;
+            font-weight: 800;
+            font-size: 1.5rem;
+            background: linear-gradient(135deg, #fff 0%, #94a3b8 100%);
+            -webkit-background-clip: text;
+            -webkit-text-fill-color: transparent;
+        }
+
+        .logo i {
+            color: var(--primary);
+            -webkit-text-fill-color: initial;
+        }
+
+        .steps-nav {
+            display: flex;
+            gap: 2rem;
+        }
+
+        .step-item {
+            display: flex;
+            align-items: center;
+            gap: 0.5rem;
+            font-size: 0.9rem;
+            font-weight: 600;
+            color: var(--text-dim);
+            transition: all 0.3s;
+            cursor: pointer;
+            padding: 0.5rem 0;
+            border-bottom: 2px solid transparent;
+        }
+
+        .step-item.active {
+            color: var(--primary);
+            border-bottom-color: var(--primary);
+        }
+
+        .step-item.completed {
+            color: var(--accent);
+        }
+
+        /* Workspace Panels */
+        .workspace {
+            flex: 1;
+            padding: 2rem;
+            display: flex;
+            align-items: center;
+            justify-content: center;
+            position: relative;
+        }
+
+        .panel {
+            width: 100%;
+            height: 100%;
+            display: none;
+            flex-direction: column;
+            align-items: center;
+            justify-content: center;
+            animation: fadeIn 0.4s ease-out;
+            position: relative;
+        }
+
+        .panel.active {
+            display: flex;
+        }
+
+        @keyframes fadeIn {
+            from { opacity: 0; transform: translateY(10px); }
+            to { opacity: 1; transform: translateY(0); }
+        }
+
+        /* Upload Zone */
+        .upload-card {
+            width: 100%;
+            max-width: 600px;
+            aspect-ratio: 4/3;
+            background: var(--surface-light);
+            border: 2px dashed var(--border);
+            border-radius: var(--radius-lg);
+            display: flex;
+            flex-direction: column;
+            align-items: center;
+            justify-content: center;
+            cursor: pointer;
+            transition: all 0.3s;
+            position: relative;
+            overflow: hidden;
+        }
+
+        .upload-card:hover {
+            border-color: var(--primary);
+            background: rgba(99, 102, 241, 0.05);
+        }
+
+        .upload-card img {
+            width: 100%;
+            height: 100%;
+            object-fit: contain;
+            display: none;
+        }
+
+        .upload-hint {
+            display: flex;
+            flex-direction: column;
+            align-items: center;
+            gap: 1rem;
+            color: var(--text-dim);
+            text-align: center;
+            padding: 2rem;
+        }
+
+        .upload-hint i {
+            width: 48px;
+            height: 48px;
+            color: var(--primary);
+        }
+
+        /* Result Viewers */
+        .viewer-wrapper {
+            width: 100%;
+            height: 100%;
+            border-radius: var(--radius-lg);
+            overflow: hidden;
+            background: #000;
+            position: relative;
+            box-shadow: 0 40px 100px rgba(0,0,0,0.6);
+        }
+
+        #frame-container {
+            width: 100%;
+            height: 100%;
+            position: relative;
+        }
+
+        .preview-frame {
+            position: absolute;
+            inset: 0;
+            width: 100%;
+            height: 100%;
+            object-fit: contain;
+            display: none;
+        }
+
+        .preview-frame.active {
+            display: block;
+        }
+
+        .viewer-overlay {
+            position: absolute;
+            bottom: 2rem;
+            left: 50%;
+            transform: translateX(-50%);
+            background: rgba(11, 15, 26, 0.6);
+            backdrop-filter: blur(12px);
+            padding: 1rem 2rem;
+            border-radius: 100px;
+            border: 1px solid var(--border);
+            display: flex;
+            align-items: center;
+            gap: 1.5rem;
+            width: 80%;
+            max-width: 600px;
+        }
+
+        /* Model Viewer Customization */
+        model-viewer {
+            width: 100%;
+            height: 100%;
+            background: radial-gradient(circle at 50% 50%, #1a2235 0%, #0b0f1a 100%);
+        }
+
+        /* Sidebar Controls */
+        .sidebar-section {
+            margin-bottom: 2rem;
+        }
+
+        .sidebar-section h3 {
+            font-size: 0.75rem;
+            text-transform: uppercase;
+            letter-spacing: 0.1em;
+            color: var(--text-dim);
+            margin-bottom: 1.25rem;
+            display: flex;
+            align-items: center;
+            gap: 0.5rem;
+        }
+
+        .control-group {
+            display: flex;
+            flex-direction: column;
+            gap: 1rem;
+        }
+
+        .input-wrapper {
+            display: flex;
+            flex-direction: column;
+            gap: 0.5rem;
+        }
+
+        .input-wrapper label {
+            font-size: 0.85rem;
+            font-weight: 600;
+            color: #cbd5e1;
+            display: flex;
+            justify-content: space-between;
+        }
+
+        .input-wrapper label span {
+            color: var(--primary);
+            font-family: monospace;
+        }
+
+        select, input[type="number"] {
+            background: var(--surface-light);
+            border: 1px solid var(--border);
+            color: white;
+            padding: 0.75rem;
+            border-radius: var(--radius-sm);
+            width: 100%;
+            outline: none;
+            transition: border-color 0.2s;
+        }
+
+        select:focus {
+            border-color: var(--primary);
+        }
+
+        input[type="range"] {
+            -webkit-appearance: none;
+            height: 4px;
+            background: var(--border);
+            border-radius: 2px;
+            margin: 10px 0;
+        }
+
+        input[type="range"]::-webkit-slider-thumb {
+            -webkit-appearance: none;
+            width: 16px;
+            height: 16px;
+            background: var(--primary);
+            border-radius: 50%;
+            cursor: pointer;
+            border: 3px solid var(--surface);
+            box-shadow: 0 0 10px rgba(129, 140, 248, 0.4);
+        }
+
+        /* Action Buttons */
+        .btn-stack {
+            margin-top: auto;
+            display: flex;
+            flex-direction: column;
+            gap: 0.75rem;
+        }
+
+        .btn {
+            width: 100%;
+            padding: 1rem;
+            border-radius: var(--radius-md);
+            font-weight: 700;
+            font-size: 0.95rem;
+            cursor: pointer;
+            transition: all 0.3s;
+            display: flex;
+            align-items: center;
+            justify-content: center;
+            gap: 0.75rem;
+            border: none;
+        }
+
+        .btn-primary {
+            background: var(--primary);
+            color: white;
+            box-shadow: 0 10px 20px rgba(99, 102, 241, 0.2);
+        }
+
+        .btn-primary:hover {
+            background: var(--primary-dark);
+            transform: translateY(-2px);
+        }
+
+        .btn-primary:disabled {
+            background: #334155;
+            color: #64748b;
+            cursor: not-allowed;
+            transform: none;
+        }
+
+        .btn-outline {
+            background: transparent;
+            border: 1px solid var(--border);
+            color: var(--text);
+        }
+
+        .btn-outline:hover {
+            background: var(--border);
+        }
+
+        /* Mode Buttons */
+        .mode-grid {
+            display: grid;
+            grid-template-columns: repeat(3, 1fr);
+            gap: 0.5rem;
+        }
+
+        .mode-tab {
+            background: var(--surface-light);
+            border: 1px solid var(--border);
+            padding: 0.5rem;
+            border-radius: var(--radius-sm);
+            font-size: 0.75rem;
+            font-weight: 600;
+            text-align: center;
+            cursor: pointer;
+            transition: all 0.2s;
+            color: var(--text-dim);
+        }
+
+        .mode-tab.active {
+            background: var(--primary);
+            color: white;
+            border-color: var(--primary);
+        }
+
+        /* Examples Footer */
+        .examples-drawer {
+            padding: 1.5rem 2rem;
+            border-top: 1px solid var(--border);
+            background: var(--surface);
+            overflow: hidden;
+            position: relative;
+        }
+
+        .examples-header {
+            display: flex;
+            align-items: center;
+            justify-content: space-between;
+            margin-bottom: 1rem;
+        }
+
+        .examples-grid {
+            display: flex;
+            flex-wrap: nowrap;
+            gap: 1rem;
+            overflow-x: auto;
+            overflow-y: hidden;
+            padding-bottom: 1rem;
+            scrollbar-width: none;
+            -ms-overflow-style: none;
+        }
+        .examples-grid::-webkit-scrollbar {
+            display: none;
+        }
+        .examples-track {
+            display: flex;
+            gap: 1rem;
+        }
+
+        .example-item {
+            flex: 0 0 140px;
+            min-width: 140px;
+            aspect-ratio: 1/1;
+            border-radius: var(--radius-md);
+            overflow: hidden;
+            cursor: pointer;
+            border: 2px solid transparent;
+            transition: all 0.2s;
+        }
+
+        .example-item:hover {
+            transform: translateY(-4px);
+            border-color: var(--primary);
+            box-shadow: 0 8px 24px rgba(99, 102, 241, 0.25);
+        }
+
+        /* Gallery Slider */
+        .gallery-slider-wrap {
+            margin-top: 0.75rem;
+            padding: 0 0.25rem;
+        }
+
+        .gallery-slider {
+            -webkit-appearance: none;
+            width: 100%;
+            height: 6px;
+            background: var(--border);
+            border-radius: 3px;
+            outline: none;
+            cursor: pointer;
+            transition: background 0.2s;
+        }
+
+        .gallery-slider:hover {
+            background: rgba(255, 255, 255, 0.12);
+        }
+
+        .gallery-slider::-webkit-slider-thumb {
+            -webkit-appearance: none;
+            width: 20px;
+            height: 20px;
+            background: linear-gradient(135deg, var(--primary), var(--accent));
+            border-radius: 50%;
+            cursor: grab;
+            box-shadow: 0 2px 8px rgba(99, 102, 241, 0.4);
+            transition: transform 0.2s, box-shadow 0.2s;
+        }
+
+        .gallery-slider::-webkit-slider-thumb:hover {
+            transform: scale(1.2);
+            box-shadow: 0 4px 16px rgba(99, 102, 241, 0.6);
+        }
+
+        .gallery-slider::-webkit-slider-thumb:active {
+            cursor: grabbing;
+            transform: scale(1.1);
+        }
+
+        .gallery-slider::-moz-range-thumb {
+            width: 20px;
+            height: 20px;
+            background: linear-gradient(135deg, var(--primary), var(--accent));
+            border-radius: 50%;
+            border: none;
+            cursor: grab;
+            box-shadow: 0 2px 8px rgba(99, 102, 241, 0.4);
+        }
+
+        .example-item img {
+            width: 100%;
+            height: 100%;
+            object-fit: cover;
+        }
+
+        /* Loading & Status */
+        .loading-overlay {
+            position: fixed;
+            inset: 0;
+            background: rgba(11, 15, 26, 0.9);
+            z-index: 1000;
+            display: none;
+            flex-direction: column;
+            align-items: center;
+            justify-content: center;
+            gap: 2rem;
+            backdrop-filter: blur(8px);
+        }
+
+        .loader-ring {
+            width: 80px;
+            height: 80px;
+            border-radius: 50%;
+            border: 4px solid var(--border);
+            border-top-color: var(--primary);
+            animation: spin 1s linear infinite;
+        }
+
+        @keyframes spin { 100% { transform: rotate(360deg); } }
+
+        .status-toast {
+            position: fixed;
+            bottom: 2rem;
+            right: 2rem;
+            background: var(--surface-light);
+            padding: 1rem 1.5rem;
+            border-radius: var(--radius-md);
+            border: 1px solid var(--border);
+            border-left: 4px solid var(--primary);
+            box-shadow: 0 20px 40px rgba(0,0,0,0.4);
+            display: none;
+            z-index: 2000;
+            animation: slideIn 0.3s cubic-bezier(0.16, 1, 0.3, 1);
+        }
+
+        @keyframes slideIn { from { transform: translateX(100%); } to { transform: translateX(0); } }
+
+        /* Scrollbar */
+        ::-webkit-scrollbar { width: 6px; height: 6px; }
+        ::-webkit-scrollbar-track { background: transparent; }
+        ::-webkit-scrollbar-thumb { background: var(--border); border-radius: 10px; }
+        ::-webkit-scrollbar-thumb:hover { background: var(--text-dim); }
+
+        /* Reference thumbnail in preview/result panels */
+        .ref-thumbnail {
+            position: absolute;
+            top: 1rem;
+            right: 1rem;
+            width: 400px;
+            height: 400px;
+            object-fit: cover;
+            border-radius: var(--radius-sm);
+            border: 2px solid var(--border);
+            z-index: 20;
+            display: none;
+            box-shadow: 0 4px 12px rgba(0,0,0,0.4);
+            cursor: pointer;
+            transition: transform 0.2s;
+        }
+
+        .ref-thumbnail:hover {
+            transform: scale(1.1);
+        }
+
+        /* Lightbox */
+        .lightbox-overlay {
+            position: fixed;
+            inset: 0;
+            background: rgba(0, 0, 0, 0.85);
+            z-index: 3000;
+            display: none;
+            align-items: center;
+            justify-content: center;
+            cursor: zoom-out;
+            backdrop-filter: blur(6px);
+        }
+
+        .lightbox-overlay img {
+            max-width: 80%;
+            max-height: 80%;
+            border-radius: var(--radius-md);
+            box-shadow: 0 20px 60px rgba(0,0,0,0.6);
+        }
+
+    </style>
+</head>
+<body>
+
+    <div class="app-shell">
+        <!-- Left Sidebar: Controls -->
+        <div class="sidebar">
+            <div class="logo" style="margin-bottom: 2.5rem;">
+                <i data-lucide="sparkles"></i>
+                <span>Pixal3D</span>
+            </div>
+
+            <div class="sidebar-section" style="margin-bottom: 1.5rem;">
+                <p style="font-size: 0.82rem; color: var(--text-dim); line-height: 1.6;">
+                    1. Upload an image and click Generate.<br>
+                    2. Click Extract GLB to export.<br>
+                    3. Download the generated GLB file.
+                </p>
+                <a href="https://ldyang694.github.io/projects/pixal3d/" target="_blank" class="btn btn-outline" style="margin-top: 1rem; padding: 0.6rem 1rem; font-size: 0.85rem;">
+                    <i data-lucide="globe" style="width: 16px;"></i>
+                    Project Page
+                </a>
+                <a href="https://huggingface.co/spaces/TencentARC/Pixal3D-Server" target="_blank" class="btn btn-outline" style="margin-top: 0.5rem; padding: 0.6rem 1rem; font-size: 0.85rem; flex-direction: column; gap: 0.25rem;">
+                    <div style="display: flex; align-items: center; gap: 0.5rem;">
+                        <i data-lucide="server" style="width: 16px;"></i>
+                        Free Deploy Server
+                    </div>
+                    <div style="display: flex; gap: 0.35rem;">
+                        <span style="background: linear-gradient(135deg, var(--accent), #059669); color: white; font-size: 0.55rem; font-weight: 700; padding: 0.1rem 0.45rem; border-radius: 4px; line-height: 1.4;">★ Recommended</span>
+                        <span style="background: linear-gradient(135deg, var(--primary), var(--primary-dark)); color: white; font-size: 0.55rem; font-weight: 700; padding: 0.1rem 0.45rem; border-radius: 4px; line-height: 1.4;">Free</span>
+                    </div>
+                </a>
+            </div>
+
+            <div class="sidebar-section">
+                <h3><i data-lucide="sliders-horizontal" style="width: 14px;"></i> Base Settings</h3>
+                <div class="control-group">
+                    <div class="input-wrapper">
+                        <label>Target Resolution</label>
+                        <select id="resolution">
+                            <option value="1024">1024 (Balanced)</option>
+                            <option value="1536">1536 (High Quality)</option>
+                        </select>
+                    </div>
+                    <div class="input-wrapper">
+                        <label>Generation Seed <span>#<span id="seed-display">42</span></span></label>
+                        <div style="display: flex; gap: 0.5rem;">
+                            <input type="number" id="seed" value="42" style="flex: 1;">
+                            <button class="btn btn-outline" style="width: 50px; padding: 0;" onclick="randomizeSeed()">
+                                <i data-lucide="rotate-cw" style="width: 16px;"></i>
+                            </button>
+                        </div>
+                    </div>
+                    <div class="input-wrapper">
+                        <p style="font-size: 0.72rem; color: var(--text-dim); line-height: 1.5; opacity: 0.7;">
+                            Note: FOV is auto-estimated via MoGe-2.<br>
+                            If distortion occurs, try manual FOV.<br>
+                            Default <b style="color: var(--primary);">0.2 rad</b> generally works well.
+                        </p>
+                    </div>
+                    <div class="input-wrapper">
+                        <label>
+                            Camera FOV
+                            <span style="display: flex; align-items: center; gap: 0.35rem;">
+                                <input type="checkbox" id="fov-auto" checked style="accent-color: var(--primary); width: 14px; height: 14px; cursor: pointer;">
+                                <span style="font-size: 0.75rem; color: var(--text-dim); cursor: pointer;" onclick="document.getElementById('fov-auto').click()">Auto</span>
+                            </span>
+                        </label>
+                        <div style="display: flex; gap: 0.5rem;">
+                            <input type="number" id="manual-fov" value="0.2" min="0.0175" max="2.9671" step="0.01" disabled style="opacity: 0.4; flex: 1;">
+                            <select id="fov-unit" disabled style="opacity: 0.4; width: 70px; padding: 0.5rem;">
+                                <option value="deg">deg</option>
+                                <option value="rad" selected>rad</option>
+                            </select>
+                        </div>
+                        <p id="fov-hint" style="font-size: 0.65rem; color: var(--text-dim); margin-top: 0.15rem;">
+                            Manual FOV in radians (0.02–2.97 rad)
+                        </p>
+                    </div>
+                </div>
+            </div>
+
+            <div class="sidebar-section" id="render-controls" style="display: none;">
+                <h3><i data-lucide="palette" style="width: 14px;"></i> Render Mode For Preview</h3>
+                <div class="mode-grid" id="mode-grid">
+                    <!-- Tabs injected via JS -->
+                </div>
+            </div>
+
+            <div class="sidebar-section">
+                <div style="display: flex; justify-content: space-between; align-items: center; margin-bottom: 1rem; cursor: pointer;" onclick="toggleAdvanced()">
+                    <h3 style="margin-bottom: 0;"><i data-lucide="shield-alert" style="width: 14px;"></i> Advanced Engine</h3>
+                    <i data-lucide="chevron-down" id="adv-chevron" style="width: 16px; transition: transform 0.3s;"></i>
+                </div>
+                <div id="advanced-settings" style="display: none; padding-top: 1rem; border-top: 1px solid var(--border);">
+                    <div class="control-group">
+                        <div class="input-wrapper">
+                            <label>SS Guidance <span><span id="ss_gs_val">7.5</span></span></label>
+                            <input type="range" id="ss_gs" min="1" max="10" step="0.1" value="7.5" oninput="updateVal('ss_gs')">
+                        </div>
+                        <div class="input-wrapper">
+                            <label>SS Sampling <span><span id="ss_steps_val">12</span></span></label>
+                            <input type="range" id="ss_steps" min="1" max="50" step="1" value="12" oninput="updateVal('ss_steps')">
+                        </div>
+                        <div class="input-wrapper">
+                            <label>Shape Guidance <span><span id="shape_gs_val">7.5</span></span></label>
+                            <input type="range" id="shape_gs" min="1" max="10" step="0.1" value="7.5" oninput="updateVal('shape_gs')">
+                        </div>
+                        <hr style="border: 0; border-top: 1px solid var(--border); margin: 0.5rem 0;">
+                        <div class="input-wrapper">
+                            <label>Decimation <span><span id="decim_val">1M</span></span></label>
+                            <input type="range" id="decimation" min="100000" max="1000000" step="10000" value="1000000" oninput="updateVal('decimation')">
+                        </div>
+                    </div>
+                </div>
+            </div>
+
+            <div class="btn-stack">
+                <button class="btn btn-primary" id="generate-btn" disabled>
+                    <i data-lucide="zap"></i>
+                    Start Generation
+                </button>
+                <button class="btn btn-outline" id="extract-btn" style="display: none;">
+                    <i data-lucide="box"></i>
+                    Extract Mesh (GLB)
+                </button>
+                <button class="btn btn-outline" id="download-btn" style="display: none; background: rgba(16, 185, 129, 0.1); border-color: var(--accent); color: var(--accent);">
+                    <i data-lucide="download"></i>
+                    Download Asset
+                </button>
+            </div>
+        </div>
+
+        <!-- Right: Main Area -->
+        <div class="main-content">
+            <header>
+                <div class="steps-nav">
+                    <div class="step-item active" id="step-1">
+                        <i data-lucide="image"></i>
+                        <span>1. SOURCE</span>
+                    </div>
+                    <div class="step-item" id="step-2">
+                        <i data-lucide="view"></i>
+                        <span>2. PREVIEW</span>
+                    </div>
+                    <div class="step-item" id="step-3">
+                        <i data-lucide="box"></i>
+                        <span>3. RESULT</span>
+                    </div>
+                </div>
+                <button class="btn btn-outline" id="clear-btn" title="Clear all & restart" style="width: 34px; height: 34px; padding: 0; border-radius: 50%; display: flex; align-items: center; justify-content: center; border-color: rgba(248,113,113,0.3);">
+                    <i data-lucide="trash-2" style="width: 16px; height: 16px; color: #f87171;"></i>
+                </button>
+            </header>
+
+            <div class="workspace">
+                <!-- Panel 1: Upload -->
+                <div class="panel active" id="panel-1">
+                    <div class="upload-card" id="drop-zone" onclick="document.getElementById('file-input').click()">
+                        <input type="file" id="file-input" hidden accept="image/*">
+                        <div class="upload-hint" id="upload-hint">
+                            <i data-lucide="cloud-upload"></i>
+                            <h2 style="font-family: 'Outfit'; margin-top: 1rem;">Upload Reference</h2>
+                            <p>Drag and drop any image, or click to browse</p>
+                        </div>
+                        <img id="source-preview" src="" alt="Source">
+                    </div>
+                </div>
+
+                <!-- Panel 2: Multi-frame Preview -->
+                <div class="panel" id="panel-2">
+                    <img id="ref-thumb-2" class="ref-thumbnail" src="" alt="Reference">
+                    <div class="viewer-wrapper">
+                        <div id="frame-container">
+                            <!-- Injected via JS -->
+                        </div>
+                        <div class="viewer-overlay">
+                            <i data-lucide="move-horizontal" style="color: var(--primary); width: 20px;"></i>
+                            <input type="range" id="angle-slider" min="0" max="7" value="0" step="1" style="flex: 1;">
+                            <div style="font-family: monospace; font-weight: 700; color: var(--primary); font-size: 0.8rem;">
+                                VIEW_ANGLE: <span id="angle-display">00</span>°
+                            </div>
+                            <div id="camera-info" style="font-family: monospace; font-size: 0.75rem; color: var(--text-dim); display: none; border-left: 1px solid var(--border); padding-left: 1rem;">
+                                <span>FOV: <span id="fov-display" style="color: var(--accent);">--</span> rad</span>
+                                <span style="margin-left: 0.75rem;">Dist: <span id="dist-display" style="color: var(--accent);">--</span></span>
+                            </div>
+                        </div>
+                    </div>
+                </div>
+
+                <!-- Panel 3: 3D Result -->
+                <div class="panel" id="panel-3">
+                    <img id="ref-thumb-3" class="ref-thumbnail" src="" alt="Reference">
+                    <div class="viewer-wrapper">
+                        <model-viewer id="main-3d-viewer" 
+                            camera-controls 
+                            auto-rotate 
+                            camera-orbit="-180deg 90deg auto"
+                            shadow-intensity="1.5" 
+                            environment-image="neutral" 
+                            exposure="1.2">
+                            <div slot="progress-bar" style="background: var(--primary); height: 4px;"></div>
+                        </model-viewer>
+                    </div>
+                </div>
+            </div>
+
+            <!-- Footer: Examples -->
+            <div class="examples-drawer">
+                <div class="examples-header">
+                    <h4 style="font-size: 0.75rem; color: var(--text-dim); text-transform: uppercase; letter-spacing: 0.1em;">Sample Gallery</h4>
+                </div>
+                <div class="examples-grid" id="examples-grid">
+                    <!-- Injected via JS -->
+                </div>
+                <div class="gallery-slider-wrap">
+                    <input type="range" id="gallery-slider" class="gallery-slider" min="0" max="100" value="0">
+                </div>
+            </div>
+        </div>
+    </div>
+
+    <div class="loading-overlay" id="loading-overlay">
+        <div class="loader-ring"></div>
+        <div style="text-align: center; width: 100%; max-width: 500px; padding: 0 2rem;">
+            <!-- Progress stages -->
+            <div id="progress-stages" style="display: none; text-align: left;">
+                <div class="progress-stage" id="progress-stage-item" style="margin-bottom: 1rem;">
+                    <div style="display: flex; justify-content: space-between; align-items: center; margin-bottom: 0.4rem;">
+                        <span id="progress-stage-name" style="font-size: 0.85rem; font-weight: 600; color: var(--primary);">Initializing...</span>
+                        <span id="progress-step-text" style="font-size: 0.75rem; color: var(--text-dim); font-family: monospace;">0/0</span>
+                    </div>
+                    <div style="width: 100%; height: 6px; background: var(--border); border-radius: 3px; overflow: hidden;">
+                        <div id="progress-bar-fill" style="width: 0%; height: 100%; background: linear-gradient(90deg, var(--primary), var(--accent)); border-radius: 3px; transition: width 0.3s ease;"></div>
+                    </div>
+                </div>
+                <!-- Stage history log -->
+                <div id="progress-log" style="font-size: 0.75rem; color: var(--text-dim); line-height: 1.8; max-height: 180px; overflow-y: auto; margin-top: 1rem; padding-top: 0.5rem; border-top: 1px solid var(--border);"></div>
+            </div>
+        </div>
+    </div>
+
+    <div class="lightbox-overlay" id="lightbox-overlay" onclick="closeLightbox()">
+        <img id="lightbox-img" src="" alt="Enlarged">
+    </div>
+
+    <div class="status-toast" id="toast">Generation started!</div>
+
+    <script type="module">
+        import { Client, handle_file } from "https://cdn.jsdelivr.net/npm/@gradio/client/dist/index.min.js";
+
+        let client;
+        let currentFile = null;
+        let preprocessedFile = null;
+        let isPreprocessing = false;
+        let generationResult = null;
+        let currentMode = "shaded_forest";
+        let currentFrame = 0;
+        const sessionId = crypto.randomUUID();
+        let currentStep = 1;
+
+        const MODES = [
+            { name: "Normal", key: "normal" },
+            { name: "Clay", key: "clay" },
+            { name: "Color", key: "base_color" },
+            { name: "Forest", key: "shaded_forest" },
+            { name: "Sunset", key: "shaded_sunset" },
+            { name: "Blue", key: "shaded_courtyard" }
+        ];
+
+        async function init() {
+            lucide.createIcons();
+            // Set default resolution based on server config (LOW_VRAM → 1024, else 1536)
+            try {
+                const cfgResp = await fetch('/app_config');
+                const cfg = await cfgResp.json();
+                document.getElementById('resolution').value = cfg.low_vram ? '1024' : '1536';
+            } catch (e) {
+                document.getElementById('resolution').value = '1536';
+            }
+            try {
+                client = await Client.connect(window.location.origin);
+                setupUI();
+                loadSamples();
+            } catch (err) {
+                console.error("Connection error:", err);
+                showToast("Connection failed. Try refreshing.");
+            }
+        }
+
+        function setupUI() {
+            // File Handling
+            const dropZone = document.getElementById('drop-zone');
+            const fileInput = document.getElementById('file-input');
+
+            dropZone.ondragover = (e) => { e.preventDefault(); dropZone.style.borderColor = 'var(--primary)'; };
+            dropZone.ondragleave = () => dropZone.style.borderColor = 'var(--border)';
+            dropZone.ondrop = (e) => {
+                e.preventDefault();
+                if (e.dataTransfer.files.length) handleImageUpload(e.dataTransfer.files[0]);
+            };
+            fileInput.onchange = (e) => { if (e.target.files.length) handleImageUpload(e.target.files[0]); };
+
+            // Buttons
+            document.getElementById('generate-btn').onclick = startGeneration;
+            document.getElementById('extract-btn').onclick = startExtraction;
+            document.getElementById('download-btn').onclick = () => {
+                const link = document.createElement('a');
+                link.href = document.getElementById('main-3d-viewer').src;
+                link.download = "pixal3d_export.glb";
+                link.click();
+            };
+
+            // Clear button
+            document.getElementById('clear-btn').onclick = () => {
+                // Reset state
+                currentFile = null;
+                generationResult = null;
+                currentFrame = 0;
+                currentMode = "shaded_forest";
+
+                // Reset source preview
+                document.getElementById('source-preview').src = '';
+                document.getElementById('source-preview').style.display = 'none';
+                document.getElementById('upload-hint').style.display = 'flex';
+                document.getElementById('file-input').value = '';
+
+                // Reset generate button
+                document.getElementById('generate-btn').disabled = true;
+
+                // Reset preview frames
+                document.getElementById('frame-container').innerHTML = '';
+                document.getElementById('angle-slider').value = 0;
+                document.getElementById('angle-display').textContent = '00';
+
+                // Reset 3D viewer
+                resetModelViewer();
+
+                // Reset thumbnails
+                document.getElementById('ref-thumb-2').style.display = 'none';
+                document.getElementById('ref-thumb-2').src = '';
+                document.getElementById('ref-thumb-3').style.display = 'none';
+                document.getElementById('ref-thumb-3').src = '';
+
+                // Reset mode tabs
+                document.querySelectorAll('.mode-tab').forEach(t => {
+                    t.classList.toggle('active', t.textContent === 'Forest');
+                });
+
+                // Go back to step 1
+                setStep(1);
+                showToast("Cleared. Ready for new upload.");
+            };
+
+            // Step navigation click
+            document.getElementById('step-1').onclick = () => setStep(1);
+            document.getElementById('step-2').onclick = () => { if (generationResult) setStep(2); };
+            document.getElementById('step-3').onclick = () => { if (document.getElementById('main-3d-viewer').src) setStep(3); };
+
+            // Slider
+            document.getElementById('angle-slider').oninput = (e) => {
+                currentFrame = parseInt(e.target.value);
+                document.getElementById('angle-display').textContent = (currentFrame * 22.5).toFixed(0).padStart(2, '0');
+                updateFrame();
+            };
+
+            // FOV auto toggle & unit switch
+            const fovAutoCheck = document.getElementById('fov-auto');
+            const fovInput = document.getElementById('manual-fov');
+            const fovUnit = document.getElementById('fov-unit');
+            const fovHint = document.getElementById('fov-hint');
+
+            function updateFovEnabled() {
+                const manual = !fovAutoCheck.checked;
+                fovInput.disabled = !manual;
+                fovUnit.disabled = !manual;
+                fovInput.style.opacity = manual ? '1' : '0.4';
+                fovUnit.style.opacity = manual ? '1' : '0.4';
+            }
+
+            fovAutoCheck.onchange = updateFovEnabled;
+
+            fovUnit.onchange = () => {
+                const val = parseFloat(fovInput.value);
+                if (isNaN(val)) return;
+                if (fovUnit.value === 'rad') {
+                    // deg -> rad
+                    fovInput.value = (val * Math.PI / 180).toFixed(4);
+                    fovInput.min = '0.0175'; fovInput.max = '2.9671'; fovInput.step = '0.01';
+                    fovHint.textContent = 'Manual FOV in radians (0.02–2.97 rad)';
+                } else {
+                    // rad -> deg
+                    fovInput.value = (val * 180 / Math.PI).toFixed(1);
+                    fovInput.min = '1'; fovInput.max = '170'; fovInput.step = '0.5';
+                    fovHint.textContent = 'Uncheck "Auto" to manually set FOV (1°–170°)';
+                }
+            };
+
+            // Mode Grid
+            const grid = document.getElementById('mode-grid');
+            MODES.forEach(m => {
+                const tab = document.createElement('div');
+                tab.className = `mode-tab ${m.key === currentMode ? 'active' : ''}`;
+                tab.textContent = m.name;
+                tab.onclick = () => {
+                    currentMode = m.key;
+                    document.querySelectorAll('.mode-tab').forEach(t => t.classList.remove('active'));
+                    tab.classList.add('active');
+                    updateFrame();
+                };
+                grid.appendChild(tab);
+            });
+        }
+
+        async function handleImageUpload(file) {
+            currentFile = file;
+            preprocessedFile = null;
+            isPreprocessing = true;
+            document.getElementById('generate-btn').disabled = true;
+            const reader = new FileReader();
+            reader.onload = (e) => {
+                const img = document.getElementById('source-preview');
+                const hint = document.getElementById('upload-hint');
+                img.src = e.target.result;
+                img.style.display = 'block';
+                hint.style.display = 'none';
+                setStep(1);
+            };
+            reader.readAsDataURL(file);
+            
+            // Call preprocess and update with segmented result
+            try {
+                const result = await client.predict("/preprocess", { image: handle_file(file) });
+                const processedUrl = result.data[0].url;
+                if (processedUrl) {
+                    preprocessedFile = processedUrl;
+                    document.getElementById('source-preview').src = processedUrl;
+                    document.getElementById('ref-thumb-2').src = processedUrl;
+                    document.getElementById('ref-thumb-2').style.display = 'block';
+                    document.getElementById('ref-thumb-3').src = processedUrl;
+                    document.getElementById('ref-thumb-3').style.display = 'block';
+                }
+            } catch (err) {
+                console.error("Preprocess failed:", err);
+            }
+            isPreprocessing = false;
+            document.getElementById('generate-btn').disabled = false;
+        }
+
+        function setStep(num) {
+            currentStep = num;
+            document.querySelectorAll('.step-item').forEach((item, i) => {
+                item.className = 'step-item';
+                if (i + 1 < num) item.classList.add('completed');
+                if (i + 1 === num) item.classList.add('active');
+            });
+            document.querySelectorAll('.panel').forEach((p, i) => {
+                p.classList.toggle('active', i + 1 === num);
+            });
+
+            // Toggle side controls based on step
+            document.getElementById('render-controls').style.display = (num >= 2) ? 'block' : 'none';
+            document.getElementById('extract-btn').style.display = (num === 2) ? 'flex' : 'none';
+            document.getElementById('download-btn').style.display = (num === 3) ? 'flex' : 'none';
+        }
+
+        async function startGeneration() {
+            if (!currentFile) return;
+
+            // Clear old preview frames and 3D result
+            generationResult = null;
+            document.getElementById('frame-container').innerHTML = '';
+            document.getElementById('angle-slider').value = 0;
+            document.getElementById('angle-display').textContent = '00';
+            resetModelViewer();
+            document.getElementById('extract-btn').style.display = 'none';
+            document.getElementById('download-btn').style.display = 'none';
+
+            // Stay on step 1 during generation
+            setStep(1);
+
+            showLoading();
+            startProgressListener();
+            try {
+                const params = {
+                    image: preprocessedFile ? handle_file(preprocessedFile) : handle_file(currentFile),
+                    seed: parseInt(document.getElementById('seed').value),
+                    resolution: parseInt(document.getElementById('resolution').value),
+                    ss_guidance_strength: parseFloat(document.getElementById('ss_gs').value),
+                    ss_sampling_steps: parseInt(document.getElementById('ss_steps').value),
+                    shape_slat_guidance_strength: parseFloat(document.getElementById('shape_gs').value),
+                    manual_fov: (() => {
+                        if (document.getElementById('fov-auto').checked) return -1.0;
+                        const v = parseFloat(document.getElementById('manual-fov').value);
+                        return document.getElementById('fov-unit').value === 'rad' ? v * 180 / Math.PI : v;
+                    })(),
+                    session_id: sessionId
+                };
+
+                const result = await client.predict("/generate_3d", params);
+                generationResult = result.data[0];
+                
+                stopProgressListener();
+                populateFrames(generationResult.render_paths);
+
+                // Display camera info
+                if (generationResult.camera_angle_x != null) {
+                    document.getElementById('fov-display').textContent = generationResult.camera_angle_x.toFixed(3);
+                    document.getElementById('dist-display').textContent = generationResult.distance.toFixed(3);
+                    document.getElementById('camera-info').style.display = 'inline';
+                }
+
+                setStep(2);
+                hideLoading();
+                showToast("Generation complete!");
+            } catch (err) {
+                console.error(err);
+                stopProgressListener();
+                hideLoading();
+                showToast("An error occurred during synthesis.");
+            }
+        }
+
+        // Progress Polling
+        let progressInterval = null;
+        let lastStageName = "";
+        
+        function startProgressListener() {
+            // Show progress UI
+            document.getElementById('progress-stages').style.display = 'block';
+            document.getElementById('progress-log').innerHTML = '';
+            document.getElementById('progress-stage-name').textContent = 'Initializing...';
+            document.getElementById('progress-step-text').textContent = '';
+            document.getElementById('progress-bar-fill').style.width = '0%';
+            lastStageName = "";
+
+            // Poll every 500ms instead of SSE
+            progressInterval = setInterval(async () => {
+                try {
+                    const resp = await fetch(`/progress?session_id=${sessionId}`);
+                    if (!resp.ok) return;
+                    const data = await resp.json();
+                    if (data.done) {
+                        stopProgressListener();
+                        return;
+                    }
+                    updateProgressUI(data);
+                } catch (e) {}
+            }, 500);
+        }
+
+        function stopProgressListener() {
+            if (progressInterval) {
+                clearInterval(progressInterval);
+                progressInterval = null;
+            }
+        }
+
+        function updateProgressUI(data) {
+            const stageName = data.stage || '';
+            const step = data.step || 0;
+            const total = data.total || 0;
+
+            // If stage changed, log the previous one as completed
+            if (stageName && stageName !== lastStageName) {
+                if (lastStageName) {
+                    const logEl = document.getElementById('progress-log');
+                    logEl.innerHTML += `<div style="display:flex;align-items:center;gap:0.4rem;"><span style="color:var(--accent);">✓</span> ${lastStageName}</div>`;
+                    logEl.scrollTop = logEl.scrollHeight;
+                }
+                lastStageName = stageName;
+            }
+
+            // Update current stage display
+            document.getElementById('progress-stage-name').textContent = stageName;
+            if (total > 0) {
+                document.getElementById('progress-step-text').textContent = `${step}/${total}`;
+                const pct = Math.min(100, (step / total) * 100);
+                document.getElementById('progress-bar-fill').style.width = pct + '%';
+            } else {
+                document.getElementById('progress-step-text').textContent = '';
+                document.getElementById('progress-bar-fill').style.width = '0%';
+            }
+        }
+
+        function populateFrames(renderPaths) {
+            const container = document.getElementById('frame-container');
+            container.innerHTML = '';
+            Object.entries(renderPaths).forEach(([mode, files]) => {
+                files.forEach((file, i) => {
+                    const img = document.createElement('img');
+                    // Try the URL from Gradio, fallback to our mounted /tmp route if it's an absolute local path
+                    let url = file.url;
+                    if (!url && file.path) {
+                        const filename = file.path.split(/[\\/]/).pop();
+                        url = `/tmp/${filename}`;
+                    }
+                    img.src = url;
+                    img.className = 'preview-frame';
+                    img.id = `frame-${mode}-${i}`;
+                    img.onerror = () => {
+                        // Fallback attempt if the first URL fails
+                        const filename = file.path ? file.path.split(/[\\/]/).pop() : null;
+                        if (filename && !img.src.includes('/tmp/')) {
+                            img.src = `/tmp/${filename}`;
+                        }
+                    };
+                    container.appendChild(img);
+                });
+            });
+            updateFrame();
+        }
+
+        function updateFrame() {
+            document.querySelectorAll('.preview-frame').forEach(f => f.classList.remove('active'));
+            const active = document.getElementById(`frame-${currentMode}-${currentFrame}`);
+            if (active) active.classList.add('active');
+        }
+
+        // Destroy and recreate model-viewer to fully purge old mesh from WebGL
+        function resetModelViewer() {
+            const container = document.querySelector('#panel-3 .viewer-wrapper');
+            const old = document.getElementById('main-3d-viewer');
+            if (old) old.remove();
+            const mv = document.createElement('model-viewer');
+            mv.id = 'main-3d-viewer';
+            mv.setAttribute('camera-controls', '');
+            mv.setAttribute('auto-rotate', '');
+            mv.setAttribute('camera-orbit', '-180deg 90deg auto');
+            mv.setAttribute('shadow-intensity', '1.5');
+            mv.setAttribute('environment-image', 'neutral');
+            mv.setAttribute('exposure', '1.2');
+            mv.style.width = '100%';
+            mv.style.height = '100%';
+            mv.style.background = 'radial-gradient(circle at 50% 50%, #1a2235 0%, #0b0f1a 100%)';
+            mv.style.visibility = 'hidden';
+            mv.innerHTML = '<div slot="progress-bar" style="background: var(--primary); height: 4px;"></div>';
+            container.appendChild(mv);
+            return mv;
+        }
+
+        async function startExtraction() {
+            if (!generationResult) return;
+
+            // Switch away from panel-3 immediately so user won't see stale mesh
+            if (currentStep === 3) setStep(2);
+
+            // Destroy old model-viewer and create a fresh one (purges WebGL scene completely)
+            const viewer = resetModelViewer();
+
+            showLoading();
+            startProgressListener();
+            try {
+                const params = {
+                    state_path: generationResult.state_path,
+                    decimation_target: parseInt(document.getElementById('decimation').value),
+                    texture_size: 4096,
+                    session_id: sessionId
+                };
+
+                const result = await client.predict("/extract_glb_api", params);
+                const glbUrl = result.data[0].url;
+
+                stopProgressListener();
+
+                // Wait for the new model to fully load before revealing (with timeout fallback)
+                await new Promise((resolve) => {
+                    let settled = false;
+                    const onLoad = () => {
+                        if (settled) return;
+                        settled = true;
+                        viewer.removeEventListener('load', onLoad);
+                        clearTimeout(timer);
+                        resolve();
+                    };
+                    // Timeout fallback: if load event never fires (e.g. model-viewer error),
+                    // resolve anyway after 30s to avoid permanent hang
+                    const timer = setTimeout(() => {
+                        if (settled) return;
+                        settled = true;
+                        viewer.removeEventListener('load', onLoad);
+                        resolve();
+                    }, 30000);
+                    viewer.addEventListener('load', onLoad);
+                    viewer.src = glbUrl;
+                });
+
+                viewer.style.visibility = 'visible';
+                setStep(3);
+                hideLoading();
+                showToast("3D Asset ready!");
+            } catch (err) {
+                console.error(err);
+                stopProgressListener();
+                viewer.style.visibility = 'visible';
+                hideLoading();
+                showToast("Extraction failed.");
+            }
+        }
+
+        function loadSamples() {
+            const grid = document.getElementById('examples-grid');
+            const samples = [
+                'assets/images/0_img.png',
+                'assets/images/1_img.png',
+                'assets/images/3_img.webp',
+                'assets/images/4_img.png',
+                'assets/images/5_img.webp',
+                'assets/images/6_img.png',
+                'assets/images/7_img.png',
+                'assets/images/9_img.png',
+                'assets/images/10_img.webp',
+                'assets/images/11_img.png',
+                'assets/images/12_img.png',
+                'assets/images/17_img.png',
+                'assets/images/21_img.png',
+                'assets/images/s_13_img.jpg',
+                'assets/images/s_14_img.jpg',
+                'assets/images/s_15_img.png',
+                'assets/images/s_16_img.png',
+                'assets/images/s_18_img.png',
+                'assets/images/s_20_img.webp',
+                'assets/images/musicman.png',
+                'assets/images/pizza.png',
+                'assets/images/sculpt.png',
+                'assets/images/treehouse.png',
+                'assets/images/warship.png',
+                'assets/images/5c80e5e03a3b60b6f03eaf555ba1dafc0e4230c472d7e8c8e2c5ca0a0dfcef10.webp',
+                'assets/images/c9340e744541f310bf89838f652602961d3e5950b31cd349bcbfc7e59e15cd2e.webp',
+                'assets/images/f94e2b76494ce2cf1874611273e5fb3d76b395793bb5647492fa85c2ce0a248b.webp'
+            ];
+
+            // Create track element
+            const track = document.createElement('div');
+            track.className = 'examples-track';
+
+            function createItem(path) {
+                const div = document.createElement('div');
+                div.className = 'example-item';
+                div.innerHTML = `<img src="${path}" draggable="false">`;
+                div.onclick = async () => {
+                    showLoading();
+                    const res = await fetch(path);
+                    const blob = await res.blob();
+                    const ext = path.split('.').pop().toLowerCase();
+                    const mimeMap = { png: 'image/png', jpg: 'image/jpeg', jpeg: 'image/jpeg', webp: 'image/webp' };
+                    const mime = mimeMap[ext] || 'image/png';
+                    const file = new File([blob], path.split('/').pop(), { type: mime });
+                    await handleImageUpload(file);
+                    hideLoading();
+                };
+                return div;
+            }
+
+            samples.forEach(path => track.appendChild(createItem(path)));
+            samples.forEach(path => track.appendChild(createItem(path)));
+
+            grid.appendChild(track);
+        }
+
+        // Helpers
+        window.openLightbox = (src) => {
+            const lb = document.getElementById('lightbox-overlay');
+            document.getElementById('lightbox-img').src = src;
+            lb.style.display = 'flex';
+        };
+
+        window.closeLightbox = () => {
+            document.getElementById('lightbox-overlay').style.display = 'none';
+        };
+
+        // Thumbnail click to enlarge
+        document.getElementById('ref-thumb-2').onclick = (e) => { e.stopPropagation(); openLightbox(e.target.src); };
+        document.getElementById('ref-thumb-3').onclick = (e) => { e.stopPropagation(); openLightbox(e.target.src); };
+
+        window.toggleAdvanced = () => {
+            const el = document.getElementById('advanced-settings');
+            const chev = document.getElementById('adv-chevron');
+            const isOpen = el.style.display === 'block';
+            el.style.display = isOpen ? 'none' : 'block';
+            chev.style.transform = isOpen ? 'rotate(0deg)' : 'rotate(180deg)';
+        };
+
+        window.updateVal = (id) => {
+            const val = document.getElementById(id).value;
+            let label = val;
+            if (id === 'decimation') label = (val/1000000).toFixed(1) + 'M';
+            document.getElementById(id + '_val').textContent = label;
+        };
+
+        window.randomizeSeed = () => {
+            const s = Math.floor(Math.random() * 999999);
+            document.getElementById('seed').value = s;
+            document.getElementById('seed-display').textContent = s;
+        };
+
+        function showLoading() {
+            document.getElementById('loading-overlay').style.display = 'flex';
+        }
+
+        function hideLoading() {
+            document.getElementById('loading-overlay').style.display = 'none';
+            document.getElementById('progress-stages').style.display = 'none';
+            document.getElementById('progress-log').innerHTML = '';
+            document.getElementById('progress-bar-fill').style.width = '0%';
+        }
+
+        function showToast(msg) {
+            const t = document.getElementById('toast');
+            t.textContent = msg;
+            t.style.display = 'block';
+            setTimeout(() => t.style.display = 'none', 3000);
+        }
+
+        // Gallery slider & scroll sync with auto-scroll
+        const exGrid = document.getElementById('examples-grid');
+        const gallerySlider = document.getElementById('gallery-slider');
+        let isSliderDragging = false;
+        let autoScrollPaused = false;
+        let pauseTimeout = null;
+        const AUTO_SCROLL_SPEED = 0.5; // px per frame
+
+        function updateSliderFromScroll() {
+            if (isSliderDragging) return;
+            const maxScroll = exGrid.scrollWidth - exGrid.clientWidth;
+            if (maxScroll <= 0) {
+                gallerySlider.value = 0;
+                return;
+            }
+            gallerySlider.value = (exGrid.scrollLeft / maxScroll) * 100;
+        }
+
+        exGrid.addEventListener('scroll', updateSliderFromScroll);
+
+        gallerySlider.addEventListener('input', () => {
+            isSliderDragging = true;
+            autoScrollPaused = true;
+            const maxScroll = exGrid.scrollWidth - exGrid.clientWidth;
+            exGrid.scrollTo({ left: (gallerySlider.value / 100) * maxScroll, behavior: 'auto' });
+        });
+        gallerySlider.addEventListener('pointerup', () => { isSliderDragging = false; resumeAutoScrollLater(); });
+        gallerySlider.addEventListener('change', () => { isSliderDragging = false; resumeAutoScrollLater(); });
+
+        // Mouse wheel horizontal scroll
+        exGrid.addEventListener('wheel', (e) => {
+            if (Math.abs(e.deltaY) > 0) {
+                e.preventDefault();
+                exGrid.scrollLeft += e.deltaY;
+                autoScrollPaused = true;
+                resumeAutoScrollLater();
+            }
+        }, { passive: false });
+
+        // Pause on hover
+        exGrid.addEventListener('mouseenter', () => { autoScrollPaused = true; });
+        exGrid.addEventListener('mouseleave', () => { resumeAutoScrollLater(500); });
+
+        function resumeAutoScrollLater(delay = 2000) {
+            clearTimeout(pauseTimeout);
+            pauseTimeout = setTimeout(() => { autoScrollPaused = false; }, delay);
+        }
+
+        // Auto-scroll loop via requestAnimationFrame
+        function autoScrollLoop() {
+            if (!autoScrollPaused) {
+                const maxScroll = exGrid.scrollWidth - exGrid.clientWidth;
+                if (maxScroll > 0) {
+                    exGrid.scrollLeft += AUTO_SCROLL_SPEED;
+                    // Loop back to start
+                    if (exGrid.scrollLeft >= maxScroll) {
+                        exGrid.scrollLeft = 0;
+                    }
+                }
+            }
+            requestAnimationFrame(autoScrollLoop);
+        }
+        requestAnimationFrame(autoScrollLoop);
+
+        init();
+    </script>
+</body>
+</html>
\ No newline at end of file
diff --git a/inference.py b/inference.py
new file mode 100644
index 0000000..4fe1c44
--- /dev/null
+++ b/inference.py
@@ -0,0 +1,312 @@
+import os
+import argparse
+import math
+import time
+import torch
+import numpy as np
+import cv2
+from PIL import Image
+
+os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1'
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+os.environ.setdefault("ATTN_BACKEND", "flash_attn")
+os.environ["FLEX_GEMM_AUTOTUNE_CACHE_PATH"] = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'autotune_cache.json')
+os.environ["FLEX_GEMM_AUTOTUNER_VERBOSE"] = '1'
+
+from pixal3d.pipelines import Pixal3DImageTo3DPipeline
+import o_voxel
+
+# ============================================================================
+# Constants & Defaults
+# ============================================================================
+
+MOGE_MODEL_NAME = "Ruicheng/moge-2-vitl"
+MODEL_PATH = "TencentARC/Pixal3D"
+
+IMAGE_COND_CONFIGS = {
+    "ss": {
+        "model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
+        "image_size": 512,
+        "grid_resolution": 16,
+    },
+    "shape_512": {
+        "model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
+        "image_size": 512,
+        "grid_resolution": 32,
+        "use_naf_upsample": True,
+        "naf_target_size": 512,
+    },
+    "shape_1024": {
+        "model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
+        "image_size": 1024,
+        "grid_resolution": 64,
+        "use_naf_upsample": True,
+        "naf_target_size": 512,
+    },
+    "tex_1024": {
+        "model_name": "camenduru/dinov3-vitl16-pretrain-lvd1689m",
+        "image_size": 1024,
+        "grid_resolution": 64,
+        "use_naf_upsample": True,
+        "naf_target_size": 1024,
+    },
+}
+
+# ============================================================================
+# Model Loading
+# ============================================================================
+
+def build_image_cond_model(config: dict):
+    from pixal3d.trainers.flow_matching.mixins.image_conditioned_proj import DinoV3ProjFeatureExtractor
+    model = DinoV3ProjFeatureExtractor(**config)
+    model.eval()
+    return model
+
+
+def load_moge_model(device="cuda", model_name=MOGE_MODEL_NAME):
+    from moge.model.v2 import MoGeModel
+    moge_model = MoGeModel.from_pretrained(model_name)
+    moge_model = moge_model.to(device)
+    moge_model.eval()
+    return moge_model
+
+
+def init_pipeline(model_path=MODEL_PATH, device="cuda", low_vram=False):
+    print(f"[Pipeline] Loading from {model_path}...")
+    pipeline = Pixal3DImageTo3DPipeline.from_pretrained(model_path)
+
+    print("[ImageCond] Building DinoV3ProjFeatureExtractor models...")
+    pipeline.image_cond_model_ss = build_image_cond_model(IMAGE_COND_CONFIGS["ss"])
+    pipeline.image_cond_model_shape_512 = build_image_cond_model(IMAGE_COND_CONFIGS["shape_512"])
+    pipeline.image_cond_model_shape_1024 = build_image_cond_model(IMAGE_COND_CONFIGS["shape_1024"])
+    pipeline.image_cond_model_tex_1024 = build_image_cond_model(IMAGE_COND_CONFIGS["tex_1024"])
+
+    if low_vram:
+        # Low-VRAM mode: models stay on CPU, loaded to GPU on-demand per stage.
+        # Peak VRAM = one flow model + one DinoV3, not all ~18 GB at once.
+        print("[NAF] Pre-downloading NAF upsampler weights (CPU only)...")
+        for attr in ['image_cond_model_ss', 'image_cond_model_shape_512',
+                     'image_cond_model_shape_1024', 'image_cond_model_tex_1024']:
+            m = getattr(pipeline, attr, None)
+            if m is not None and getattr(m, 'use_naf_upsample', False):
+                m._load_naf()
+        pipeline._device = torch.device(device)
+        pipeline.low_vram = True
+        print("[Pipeline] Low-VRAM mode enabled.")
+    else:
+        # Standard mode: all models loaded to GPU at once (faster, needs more VRAM).
+        pipeline.low_vram = False
+        pipeline.cuda()
+        pipeline.image_cond_model_ss.cuda()
+        pipeline.image_cond_model_shape_512.cuda()
+        pipeline.image_cond_model_shape_1024.cuda()
+        pipeline.image_cond_model_tex_1024.cuda()
+        print("[NAF] Pre-loading NAF upsampler model...")
+        for attr in ['image_cond_model_ss', 'image_cond_model_shape_512',
+                     'image_cond_model_shape_1024', 'image_cond_model_tex_1024']:
+            m = getattr(pipeline, attr, None)
+            if m is not None and getattr(m, 'use_naf_upsample', False):
+                m._load_naf()
+        print("[Pipeline] Standard mode (all models on GPU).")
+
+    return pipeline
+
+# ============================================================================
+# Camera Estimation
+# ============================================================================
+
+def compute_f_pixels(camera_angle_x: float, resolution: int) -> float:
+    focal_length = 16.0 / torch.tan(torch.tensor(camera_angle_x / 2.0))
+    f_pixels = focal_length * resolution / 32.0
+    return float(f_pixels.item())
+
+
+def distance_from_fov(camera_angle_x, grid_point, target_point, mesh_scale, image_resolution):
+    rotation_matrix = torch.tensor([[1.0, 0.0, 0.0], [0.0, 0.0, -1.0], [0.0, 1.0, 0.0]])
+    gp = grid_point.to(torch.float32) @ rotation_matrix.T
+    gp = gp / mesh_scale / 2
+    xw, yw, zw = gp[0].item(), gp[1].item(), gp[2].item()
+    xt, yt = float(target_point[0].item()), float(target_point[1].item())
+    f_pixels = compute_f_pixels(camera_angle_x, image_resolution)
+    x_ndc = xt - image_resolution / 2.0
+    y_ndc = -(yt - image_resolution / 2.0)
+    distance_x = f_pixels * xw / x_ndc - yw
+    return {"distance_from_x": float(distance_x), "f_pixels": float(f_pixels)}
+
+
+def get_camera_params_wild_moge(image_path, moge_model, device="cuda", mesh_scale=1.0, extend_pixel=0, image_resolution=512):
+    pil_image = Image.open(image_path).convert("RGB")
+    width, height = pil_image.size
+    image_np = np.array(pil_image).astype(np.float32) / 255.0
+    image_tensor = torch.from_numpy(image_np).permute(2, 0, 1).to(device)
+    with torch.no_grad():
+        output = moge_model.infer(image_tensor)
+    intrinsics = output["intrinsics"].squeeze().cpu().numpy()
+    fx_normalized = intrinsics[0, 0]
+    fx = fx_normalized * width
+    camera_angle_x = 2 * math.atan(width / (2 * fx))
+
+    grid_point = torch.tensor([-1.0, 0.0, 0.0])
+    distance = distance_from_fov(
+        camera_angle_x, grid_point,
+        torch.tensor([0 - extend_pixel, image_resolution - 1 + extend_pixel]),
+        mesh_scale, image_resolution
+    )["distance_from_x"]
+    return {'camera_angle_x': camera_angle_x, 'distance': distance, 'mesh_scale': mesh_scale}
+
+# ============================================================================
+# Main Inference
+# ============================================================================
+
+def run_inference(
+    image_path: str,
+    output_path: str,
+    seed: int = 42,
+    ss_guidance_strength: float = 7.5,
+    ss_guidance_rescale: float = 0.7,
+    ss_sampling_steps: int = 12,
+    ss_rescale_t: float = 5.0,
+    shape_slat_guidance_strength: float = 7.5,
+    shape_slat_guidance_rescale: float = 0.5,
+    shape_slat_sampling_steps: int = 12,
+    shape_slat_rescale_t: float = 3.0,
+    tex_slat_guidance_strength: float = 1.0,
+    tex_slat_guidance_rescale: float = 0.0,
+    tex_slat_sampling_steps: int = 12,
+    tex_slat_rescale_t: float = 3.0,
+    mesh_scale: float = 1.0,
+    extend_pixel: int = 0,
+    image_resolution: int = 512,
+    max_num_tokens: int = 49152,
+    model_path: str = MODEL_PATH,
+    manual_fov: float = -1.0,
+    low_vram: bool = False,
+    resolution: int = -1,
+):
+    # Load models
+    pipeline = init_pipeline(model_path, low_vram=low_vram)
+
+    # Preprocess image first — rembg loads to GPU for this call, then offloads.
+    # MoGe is loaded afterwards so both never occupy VRAM at the same time.
+    print(f"[Inference] Processing image: {image_path}")
+    img = Image.open(image_path)
+    image_preprocessed = pipeline.preprocess_image(img)
+
+    # Save preprocessed image for MoGe
+    tmp_path = os.path.join(os.path.dirname(os.path.abspath(output_path)), f"_tmp_preprocessed_{int(time.time()*1000)}.png")
+    image_preprocessed.save(tmp_path)
+
+    # Camera estimation
+    if manual_fov > 0:
+        # Use manually specified FOV (in radians)
+        camera_angle_x = float(manual_fov)
+        grid_point = torch.tensor([-1.0, 0.0, 0.0])
+        distance = distance_from_fov(
+            camera_angle_x, grid_point,
+            torch.tensor([0 - extend_pixel, image_resolution - 1 + extend_pixel]),
+            mesh_scale, image_resolution
+        )["distance_from_x"]
+        camera_params = {'camera_angle_x': camera_angle_x, 'distance': distance, 'mesh_scale': mesh_scale}
+        print(f"[Inference] Using manual FOV: {math.degrees(manual_fov):.2f}° ({manual_fov:.4f} rad), distance={distance:.4f}")
+    else:
+        print("[MoGe-2] Loading model for camera estimation...")
+        moge_model = load_moge_model(device="cuda")
+        print("[Inference] Estimating camera parameters...")
+        camera_params = get_camera_params_wild_moge(
+            tmp_path, moge_model, device="cuda",
+            mesh_scale=mesh_scale, extend_pixel=extend_pixel,
+            image_resolution=image_resolution,
+        )
+        print(f"  camera_angle_x={camera_params['camera_angle_x']:.4f}, distance={camera_params['distance']:.4f}")
+        # MoGe is only needed for camera estimation; free its VRAM for inference.
+        moge_model.cpu()
+        del moge_model
+        torch.cuda.empty_cache()
+    os.remove(tmp_path)
+
+    # Run pipeline
+    print("[Inference] Running 3D generation pipeline...")
+    torch.manual_seed(seed)
+
+    ss_sampler_override = {
+        "steps": ss_sampling_steps, "guidance_strength": ss_guidance_strength,
+        "guidance_rescale": ss_guidance_rescale, "rescale_t": ss_rescale_t,
+    }
+    shape_sampler_override = {
+        "steps": shape_slat_sampling_steps, "guidance_strength": shape_slat_guidance_strength,
+        "guidance_rescale": shape_slat_guidance_rescale, "rescale_t": shape_slat_rescale_t,
+    }
+    tex_sampler_override = {
+        "steps": tex_slat_sampling_steps, "guidance_strength": tex_slat_guidance_strength,
+        "guidance_rescale": tex_slat_guidance_rescale, "rescale_t": tex_slat_rescale_t,
+    }
+
+    pipeline_type = f"{resolution if resolution > 0 else (1024 if low_vram else 1536)}_cascade"
+    print(f"[Inference] Using pipeline_type={pipeline_type}")
+    mesh_list, (shape_slat, tex_slat, res) = pipeline.run(
+        image_preprocessed,
+        camera_params=camera_params,
+        seed=seed,
+        sparse_structure_sampler_params=ss_sampler_override,
+        shape_slat_sampler_params=shape_sampler_override,
+        tex_slat_sampler_params=tex_sampler_override,
+        preprocess_image=False,
+        return_latent=True,
+        pipeline_type=pipeline_type,
+        max_num_tokens=max_num_tokens,
+    )
+
+    mesh = mesh_list[0]
+
+    # Extract GLB
+    print("[Inference] Extracting GLB...")
+    glb = o_voxel.postprocess.to_glb(
+        vertices=mesh.vertices, faces=mesh.faces, attr_volume=mesh.attrs,
+        coords=mesh.coords, attr_layout=pipeline.pbr_attr_layout,
+        grid_size=res, aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
+        decimation_target=1000000, texture_size=4096,
+        remesh=True, remesh_band=1, remesh_project=0, use_tqdm=True,
+    )
+
+    # Apply rotation
+    rot = np.array([
+        [-1,  0,  0,  0],
+        [ 0,  0, -1,  0],
+        [ 0, -1,  0,  0],
+        [ 0,  0,  0,  1],
+    ], dtype=np.float64)
+    glb.apply_transform(rot)
+
+    # Export
+    os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
+    glb.export(output_path, extension_webp=True)
+    print(f"[Done] GLB saved to: {output_path}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Pixal3D Inference: Image to GLB")
+    parser.add_argument("--image", type=str, required=True, help="Path to input image")
+    parser.add_argument("--output", type=str, default="./output.glb", help="Output GLB file path")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed")
+    parser.add_argument("--fov", type=float, default=-1.0,
+                        help="Manual camera FOV in radians (e.g. 0.2). "
+                             "If not set, FOV is auto-estimated via MoGe-2. "
+                             "Try 0.2 rad if you notice distortion.")
+    parser.add_argument("--model_path", type=str, default=MODEL_PATH, help="Model path or HuggingFace repo")
+    parser.add_argument("--low_vram", action="store_true",
+                        help="Enable low-VRAM mode: models stay on CPU and are loaded to GPU on-demand per stage. "
+                             "Reduces peak VRAM from ~18GB to ~10-12GB at the cost of slower inference.")
+    parser.add_argument("--resolution", type=int, default=-1,
+                        help="Pipeline resolution (1024 or 1536). Default: 1024 if --low_vram, else 1536.")
+
+    args = parser.parse_args()
+
+    run_inference(
+        image_path=args.image,
+        output_path=args.output,
+        seed=args.seed,
+        manual_fov=args.fov,
+        model_path=args.model_path,
+        low_vram=args.low_vram,
+        resolution=args.resolution,
+    )
diff --git a/pixal3d/__init__.py b/pixal3d/__init__.py
new file mode 100644
index 0000000..20d240a
--- /dev/null
+++ b/pixal3d/__init__.py
@@ -0,0 +1,6 @@
+from . import models
+from . import modules
+from . import pipelines
+from . import renderers
+from . import representations
+from . import utils
diff --git a/pixal3d/datasets/__init__.py b/pixal3d/datasets/__init__.py
new file mode 100644
index 0000000..eadadad
--- /dev/null
+++ b/pixal3d/datasets/__init__.py
@@ -0,0 +1,52 @@
+import importlib
+
+__attributes = {
+    'FlexiDualGridDataset': 'flexi_dual_grid',
+    'SparseVoxelPbrDataset':'sparse_voxel_pbr',
+    
+    'SparseStructureLatent': 'sparse_structure_latent',
+    'TextConditionedSparseStructureLatent': 'sparse_structure_latent',
+    'ImageConditionedSparseStructureLatent': 'sparse_structure_latent',
+    'SparseStructureLatentView': 'sparse_structure_latent',
+    'ViewImageConditionedSparseStructureLatentView': 'sparse_structure_latent',
+    
+    'SLat': 'structured_latent',
+    'ImageConditionedSLat': 'structured_latent',
+    'SLatShape': 'structured_latent_shape',
+    'ImageConditionedSLatShape': 'structured_latent_shape',
+    'SLatShapeView': 'structured_latent_shape',
+    'ViewImageConditionedSLatShapeView': 'structured_latent_shape',
+    'SLatPbr': 'structured_latent_svpbr',
+    'ImageConditionedSLatPbr': 'structured_latent_svpbr',
+    'SLatPbrView': 'structured_latent_svpbr',
+    'ViewImageConditionedSLatPbrView': 'structured_latent_svpbr',
+}
+
+__submodules = []
+
+__all__ = list(__attributes.keys()) + __submodules
+
+def __getattr__(name):
+    if name not in globals():
+        if name in __attributes:
+            module_name = __attributes[name]
+            module = importlib.import_module(f".{module_name}", __name__)
+            globals()[name] = getattr(module, name)
+        elif name in __submodules:
+            module = importlib.import_module(f".{name}", __name__)
+            globals()[name] = module
+        else:
+            raise AttributeError(f"module {__name__} has no attribute {name}")
+    return globals()[name]
+
+
+# For Pylance
+if __name__ == '__main__':    
+    from .flexi_dual_grid import FlexiDualGridDataset
+    from .sparse_voxel_pbr import SparseVoxelPbrDataset
+    
+    from .sparse_structure_latent import SparseStructureLatent, ImageConditionedSparseStructureLatent
+    from .structured_latent import SLat, ImageConditionedSLat
+    from .structured_latent_shape import SLatShape, ImageConditionedSLatShape
+    from .structured_latent_svpbr import SLatPbr, ImageConditionedSLatPbr
+    
\ No newline at end of file
diff --git a/pixal3d/datasets/components.py b/pixal3d/datasets/components.py
new file mode 100644
index 0000000..af9090a
--- /dev/null
+++ b/pixal3d/datasets/components.py
@@ -0,0 +1,349 @@
+from typing import *
+import json
+from abc import abstractmethod
+import os
+import json
+import torch
+import numpy as np
+import pandas as pd
+from PIL import Image
+from torch.utils.data import Dataset
+
+
+class StandardDatasetBase(Dataset):
+    """
+    Base class for standard datasets.
+
+    Args:
+        roots (str): paths to the dataset
+        skip_list (str, optional): path to a file containing sha256 hashes to skip (one per line)
+                                   Format: "dataset/sha256" (e.g., "ABO/6a79dbb5...")
+        skip_aesthetic_score_datasets (list, optional): list of dataset names to skip aesthetic score check
+                                                        (e.g., ["texverse"] for datasets without aesthetic_score)
+    """
+
+    def __init__(self,
+        roots: str,
+        skip_list: Optional[str] = None,
+        skip_aesthetic_score_datasets: Optional[List[str]] = None,
+    ):
+        super().__init__()
+        
+        # Datasets to skip aesthetic score check
+        self.skip_aesthetic_score_datasets = set(skip_aesthetic_score_datasets or [])
+        
+        # Load skip list if provided
+        self.skip_set = set()
+        if skip_list is not None and os.path.exists(skip_list):
+            with open(skip_list, 'r') as f:
+                for line in f:
+                    line = line.strip()
+                    if line and not line.startswith('#'):
+                        self.skip_set.add(line)
+            print(f'Loaded {len(self.skip_set)} items from skip_list: {skip_list}')
+        
+        try:
+            self.roots = json.loads(roots)
+            root_type = 'obj'
+        except:
+            self.roots = roots.split(',')
+            root_type = 'list'
+        self.instances = []
+        self.metadata = pd.DataFrame()
+        
+        self._stats = {}
+        if root_type == 'obj':
+            for key, root in self.roots.items():
+                self._stats[key] = {}
+                metadata = pd.DataFrame(columns=['sha256']).set_index('sha256')
+                
+                # Only merge key fields from ss_latent and render_cond
+                # Exclude base, because cond_rendered=False in base/metadata.csv would incorrectly overwrite real values
+                for sub_key, r in root.items():
+                    if sub_key == 'base':
+                        continue  # Skip base directory
+                    metadata_file = os.path.join(r, 'metadata.csv')
+                    if os.path.exists(metadata_file):
+                        metadata = metadata.combine_first(pd.read_csv(metadata_file).set_index('sha256'))
+                
+                # Read aesthetic_score separately from base (avoid reading other potentially conflicting columns)
+                if 'base' in root:
+                    base_metadata_file = os.path.join(root['base'], 'metadata.csv')
+                    if os.path.exists(base_metadata_file):
+                        base_df = pd.read_csv(base_metadata_file).set_index('sha256')
+                        if 'aesthetic_score' in base_df.columns and 'aesthetic_score' not in metadata.columns:
+                            metadata['aesthetic_score'] = base_df['aesthetic_score']
+                
+                self._stats[key]['Total'] = len(metadata)
+                metadata, stats = self.filter_metadata(metadata, dataset_name=key)
+                self._stats[key].update(stats)
+                
+                # Filter out items in skip_list
+                skipped_count = 0
+                for sha256 in metadata.index.values:
+                    skip_key = f'{key}/{sha256}'
+                    if skip_key in self.skip_set:
+                        skipped_count += 1
+                    else:
+                        self.instances.append((root, sha256, key))
+                if skipped_count > 0:
+                    self._stats[key]['Skipped (skip_list)'] = skipped_count
+                    self._stats[key]['After skip_list'] = len(metadata) - skipped_count
+                
+                self.metadata = pd.concat([self.metadata, metadata])
+        else:
+            for root in self.roots:
+                key = os.path.basename(root)
+                self._stats[key] = {}
+                metadata = pd.read_csv(os.path.join(root, 'metadata.csv'))
+                self._stats[key]['Total'] = len(metadata)
+                metadata, stats = self.filter_metadata(metadata, dataset_name=key)
+                self._stats[key].update(stats)
+                
+                # Filter out items in skip_list
+                skipped_count = 0
+                for sha256 in metadata['sha256'].values:
+                    skip_key = f'{key}/{sha256}'
+                    if skip_key in self.skip_set:
+                        skipped_count += 1
+                    else:
+                        self.instances.append((root, sha256, key))
+                if skipped_count > 0:
+                    self._stats[key]['Skipped (skip_list)'] = skipped_count
+                    self._stats[key]['After skip_list'] = len(metadata) - skipped_count
+                metadata.set_index('sha256', inplace=True)
+                self.metadata = pd.concat([self.metadata, metadata])
+            
+    @abstractmethod
+    def filter_metadata(self, metadata: pd.DataFrame, dataset_name: str = None) -> Tuple[pd.DataFrame, Dict[str, int]]:
+        pass
+    
+    @abstractmethod
+    def get_instance(self, root, instance: str) -> Dict[str, Any]:
+        pass
+        
+    def __len__(self):
+        return len(self.instances)
+
+    def __getitem__(self, index) -> Dict[str, Any]:
+        try:
+            root, instance, dataset_name = self.instances[index]
+            pack = self.get_instance(root, instance)
+            pack['_dataset_name'] = dataset_name
+            pack['_sha256'] = instance
+            return pack
+        except Exception as e:
+            print(f'Error loading {self.instances[index][1]}: {e}')
+            return self.__getitem__(np.random.randint(0, len(self)))
+        
+    def __str__(self):
+        lines = []
+        lines.append(self.__class__.__name__)
+        lines.append(f'  - Total instances: {len(self)}')
+        lines.append(f'  - Sources:')
+        for key, stats in self._stats.items():
+            lines.append(f'    - {key}:')
+            for k, v in stats.items():
+                lines.append(f'      - {k}: {v}')
+        return '\n'.join(lines)
+
+
+class ImageConditionedMixin:
+    def __init__(self, roots, *, image_size=518, **kwargs):
+        self.image_size = image_size
+        super().__init__(roots, **kwargs)
+    
+    def filter_metadata(self, metadata, dataset_name=None):
+        metadata, stats = super().filter_metadata(metadata, dataset_name=dataset_name)
+        metadata = metadata[metadata['cond_rendered'].notna()]
+        stats['Cond rendered'] = len(metadata)
+        return metadata, stats
+    
+    def get_instance(self, root, instance):
+        pack = super().get_instance(root, instance)
+       
+        image_root = os.path.join(root['render_cond'], instance)
+        with open(os.path.join(image_root, 'transforms.json')) as f:
+            metadata = json.load(f)
+        n_views = len(metadata['frames'])
+        view = np.random.randint(n_views)
+        metadata = metadata['frames'][view]
+
+        image_path = os.path.join(image_root, metadata['file_path'])
+        image = Image.open(image_path)
+
+        alpha = np.array(image.getchannel(3))
+        bbox = np.array(alpha).nonzero()
+        bbox = [bbox[1].min(), bbox[0].min(), bbox[1].max(), bbox[0].max()]
+        center = [(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2]
+        hsize = max(bbox[2] - bbox[0], bbox[3] - bbox[1]) / 2
+        aug_hsize = hsize
+        aug_center_offset = [0, 0]
+        aug_center = [center[0] + aug_center_offset[0], center[1] + aug_center_offset[1]]
+        aug_bbox = [int(aug_center[0] - aug_hsize), int(aug_center[1] - aug_hsize), int(aug_center[0] + aug_hsize), int(aug_center[1] + aug_hsize)]
+        image = image.crop(aug_bbox)
+
+        image = image.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS)
+        alpha = image.getchannel(3)
+        image = image.convert('RGB')
+        image = torch.tensor(np.array(image)).permute(2, 0, 1).float() / 255.0
+        alpha = torch.tensor(np.array(alpha)).float() / 255.0
+        image = image * alpha.unsqueeze(0)
+        pack['cond'] = image
+       
+        return pack
+
+
+class ViewImageConditionedMixin:
+    """
+    Mixin for view-based image-conditioned datasets.
+    
+    This mixin is designed for datasets where ss_latent is stored per-view (view{XX}.npz),
+    and needs to load the corresponding view image and scale from view{XX}_scale.json.
+    
+    Args:
+        image_size: Target image size
+        load_camera_info: Whether to load camera information for view-aligned conditioning
+    """
+    def __init__(self, roots, *, image_size=518, load_camera_info=False, **kwargs):
+        self.image_size = image_size
+        # self.load_camera_info = load_camera_info
+        super().__init__(roots, **kwargs)
+    
+    def filter_metadata(self, metadata, dataset_name=None):
+        metadata, stats = super().filter_metadata(metadata, dataset_name=dataset_name)
+        metadata = metadata[metadata['cond_rendered'].notna()]
+        stats['Cond rendered'] = len(metadata)
+        return metadata, stats
+    
+    def get_instance(self, root, instance):
+        """
+        Get instance with view-aligned image and camera info.
+        
+        Expects parent class to set:
+            - pack['x_0']: the latent tensor
+            - self._current_view_idx: the selected view index
+            - self._current_latent_dir: the latent directory path
+        """
+        pack = super().get_instance(root, instance)
+        
+        # Get view_idx from parent class (set by SparseStructureLatentView)
+        if not hasattr(self, '_current_view_idx'):
+            raise RuntimeError("Parent class must set '_current_view_idx' before calling ViewImageConditionedMixin.get_instance")
+        if not hasattr(self, '_current_latent_dir'):
+            raise RuntimeError("Parent class must set '_current_latent_dir' before calling ViewImageConditionedMixin.get_instance")
+        view_idx = self._current_view_idx
+        latent_dir = self._current_latent_dir
+        
+        # Load image metadata
+        image_root = os.path.join(root['render_cond'], instance)
+        with open(os.path.join(image_root, 'transforms.json')) as f:
+            metadata = json.load(f)
+        
+        # Load corresponding image for this view
+        frame_metadata = metadata['frames'][view_idx]
+        image_path = os.path.join(image_root, frame_metadata['file_path'])
+        image = Image.open(image_path)
+
+        image = image.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS)
+        alpha = image.getchannel(3)
+        image = image.convert('RGB')
+        image = torch.tensor(np.array(image)).permute(2, 0, 1).float() / 255.0
+        alpha = torch.tensor(np.array(alpha)).float() / 255.0
+        image = image * alpha.unsqueeze(0)
+        pack['cond'] = image
+        
+        # Load camera info if requested
+   
+        # camera_angle_x: check frame first, then root metadata
+        if 'camera_angle_x' in frame_metadata:
+            camera_angle_x = float(frame_metadata['camera_angle_x'])
+        elif 'camera_angle_x' in metadata:
+            camera_angle_x = float(metadata['camera_angle_x'])
+        else:
+            raise KeyError(f"'camera_angle_x' not found in transforms.json for {instance}")
+        pack['camera_angle_x'] = torch.tensor(camera_angle_x, dtype=torch.float32)
+        
+        # transform_matrix
+        if 'transform_matrix' not in frame_metadata:
+            raise KeyError(f"'transform_matrix' not found in frame {view_idx} for {instance}")
+        transform_matrix = torch.tensor(frame_metadata['transform_matrix'], dtype=torch.float32)
+        distance = torch.norm(transform_matrix[:3, 3]).item()
+            
+        pack['camera_distance'] = torch.tensor(distance, dtype=torch.float32)
+        # NOTE: Do NOT pass transform_matrix to ProjGrid.
+        # shape_latent space objects are already rotated to front-view by transform_mesh,
+        # so ProjGrid should use the default front_view_transform_matrix + distance.
+        # pack['transform_matrix'] = transform_matrix
+        
+        # Load mesh_scale from ss_latent directory's view{XX}_scale.json
+        scale_json_path = os.path.join(latent_dir, f'view{view_idx:02d}_scale.json')
+        if not os.path.exists(scale_json_path):
+            raise FileNotFoundError(f"Scale file not found: {scale_json_path}")
+        with open(scale_json_path) as f:
+            scale_data = json.load(f)
+        if 'total_scale' not in scale_data:
+            raise KeyError(f"'total_scale' not found in {scale_json_path}")
+        pack['mesh_scale'] = torch.tensor(float(scale_data['total_scale']), dtype=torch.float32)
+       
+        return pack
+
+
+class MultiImageConditionedMixin:
+    def __init__(self, roots, *, image_size=518, max_image_cond_view = 4, **kwargs):
+        self.image_size = image_size
+        self.max_image_cond_view = max_image_cond_view
+        super().__init__(roots, **kwargs)
+
+    def filter_metadata(self, metadata, dataset_name=None):
+        metadata, stats = super().filter_metadata(metadata, dataset_name=dataset_name)
+        metadata = metadata[metadata['cond_rendered'].notna()]
+        stats['Cond rendered'] = len(metadata)
+        return metadata, stats
+    
+    def get_instance(self, root, instance):
+        pack = super().get_instance(root, instance)
+       
+        image_root = os.path.join(root['render_cond'], instance)
+        with open(os.path.join(image_root, 'transforms.json')) as f:
+            metadata = json.load(f)
+
+        n_views = len(metadata['frames'])
+        n_sample_views = np.random.randint(1, self.max_image_cond_view+1)
+
+        assert n_views >= n_sample_views, f'Not enough views to sample {n_sample_views} unique images.'
+
+        sampled_views = np.random.choice(n_views, size=n_sample_views, replace=False)
+
+        cond_images = []
+        for v in sampled_views:
+            frame_info = metadata['frames'][v]
+            image_path = os.path.join(image_root, frame_info['file_path'])
+            image = Image.open(image_path)
+
+            alpha = np.array(image.getchannel(3))
+            bbox = np.array(alpha).nonzero()
+            bbox = [bbox[1].min(), bbox[0].min(), bbox[1].max(), bbox[0].max()]
+            center = [(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2]
+            hsize = max(bbox[2] - bbox[0], bbox[3] - bbox[1]) / 2
+            aug_hsize = hsize
+            aug_center = center
+            aug_bbox = [
+                int(aug_center[0] - aug_hsize),
+                int(aug_center[1] - aug_hsize),
+                int(aug_center[0] + aug_hsize),
+                int(aug_center[1] + aug_hsize),
+            ]
+
+            img = image.crop(aug_bbox)
+            img = img.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS)
+            alpha = img.getchannel(3)
+            img = img.convert('RGB')
+            img = torch.tensor(np.array(img)).permute(2, 0, 1).float() / 255.0
+            alpha = torch.tensor(np.array(alpha)).float() / 255.0
+            img = img * alpha.unsqueeze(0)
+
+            cond_images.append(img)
+
+        pack['cond'] = [torch.stack(cond_images, dim=0)]  # (V,3,H,W)
+        return pack
diff --git a/pixal3d/datasets/flexi_dual_grid.py b/pixal3d/datasets/flexi_dual_grid.py
new file mode 100644
index 0000000..f3a9353
--- /dev/null
+++ b/pixal3d/datasets/flexi_dual_grid.py
@@ -0,0 +1,173 @@
+import os
+import numpy as np
+import pickle
+import torch
+import utils3d
+from .components import StandardDatasetBase
+from ..modules import sparse as sp
+from ..renderers import MeshRenderer
+from ..representations import Mesh
+from ..utils.data_utils import load_balanced_group_indices
+import o_voxel
+
+
+class FlexiDualGridVisMixin:
+    @torch.no_grad()
+    def visualize_sample(self, x: dict):
+        mesh = x['mesh']
+        
+        renderer = MeshRenderer({'near': 1, 'far': 3})
+        renderer.rendering_options.resolution = 512
+        renderer.rendering_options.ssaa = 4
+        
+        # Build camera
+        yaws = [0, np.pi / 2, np.pi, 3 * np.pi / 2]
+        yaws_offset = np.random.uniform(-np.pi / 4, np.pi / 4)
+        yaws = [y + yaws_offset for y in yaws]
+        pitch = [np.random.uniform(-np.pi / 4, np.pi / 4) for _ in range(4)]
+
+        exts = []
+        ints = []
+        for yaw, pitch in zip(yaws, pitch):
+            orig = torch.tensor([
+                np.sin(yaw) * np.cos(pitch),
+                np.cos(yaw) * np.cos(pitch),
+                np.sin(pitch),
+            ]).float().cuda() * 2
+            fov = torch.deg2rad(torch.tensor(30)).cuda()
+            extrinsics = utils3d.torch.extrinsics_look_at(orig, torch.tensor([0, 0, 0]).float().cuda(), torch.tensor([0, 0, 1]).float().cuda())
+            intrinsics = utils3d.torch.intrinsics_from_fov_xy(fov, fov)
+            exts.append(extrinsics)
+            ints.append(intrinsics)
+        
+        # Build each representation
+        images = []
+        for m in mesh:
+            image = torch.zeros(3, 1024, 1024).cuda()
+            tile = [2, 2]
+            for j, (ext, intr) in enumerate(zip(exts, ints)):
+                image[:, 512 * (j // tile[1]):512 * (j // tile[1] + 1), 512 * (j % tile[1]):512 * (j % tile[1] + 1)] = \
+                    renderer.render(m.cuda(), ext, intr)['normal']
+            images.append(image)
+        images = torch.stack(images)
+        
+        return images
+    
+
+class FlexiDualGridDataset(FlexiDualGridVisMixin, StandardDatasetBase):
+    """
+    Flexible Dual Grid Dataset
+    
+    Args:
+        roots (str): path to the dataset
+        resolution (int): resolution of the voxel grid
+        min_aesthetic_score (float): minimum aesthetic score of the instances to be included in the dataset
+    """
+
+    def __init__(
+        self,
+        roots,
+        resolution: int = 1024,
+        max_active_voxels: int = 1000000,
+        max_num_faces: int = None,
+        min_aesthetic_score: float = 5.0,
+    ):
+        self.resolution = resolution
+        self.min_aesthetic_score = min_aesthetic_score
+        self.max_active_voxels = max_active_voxels
+        self.max_num_faces = max_num_faces
+        self.value_range = (0, 1)
+
+        super().__init__(roots)
+        
+        self.loads = [self.metadata.loc[sha256, f'dual_grid_size'] for _, sha256, _ in self.instances]
+        
+    def __str__(self):
+        lines = [
+            super().__str__(),
+            f'  - Resolution: {self.resolution}',
+        ]
+        return '\n'.join(lines)
+        
+    def filter_metadata(self, metadata, dataset_name=None):
+        stats = {}
+        metadata = metadata[metadata[f'dual_grid_converted'] == True]
+        stats['Dual Grid Converted'] = len(metadata)
+        if self.min_aesthetic_score is not None:
+            metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]
+            stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)
+        metadata = metadata[metadata[f'dual_grid_size'] <= self.max_active_voxels]
+        stats[f'Active Voxels <= {self.max_active_voxels}'] = len(metadata)
+        if self.max_num_faces is not None:
+            metadata = metadata[metadata['num_faces'] <= self.max_num_faces]
+            stats[f'Faces <= {self.max_num_faces}'] = len(metadata)
+        return metadata, stats
+    
+    def read_mesh(self, root, instance):
+        with open(os.path.join(root, f'{instance}.pickle'), 'rb') as f:
+            dump = pickle.load(f)
+        start = 0
+        vertices = []
+        faces = []
+        for obj in dump['objects']:
+            if obj['vertices'].size == 0 or obj['faces'].size == 0:
+                continue
+            vertices.append(obj['vertices'])
+            faces.append(obj['faces'] + start)
+            start += len(obj['vertices'])
+        vertices = torch.from_numpy(np.concatenate(vertices, axis=0)).float()
+        faces = torch.from_numpy(np.concatenate(faces, axis=0)).long()
+        vertices_min = vertices.min(dim=0)[0]
+        vertices_max = vertices.max(dim=0)[0]
+        center = (vertices_min + vertices_max) / 2
+        scale = 0.99999 / (vertices_max - vertices_min).max()
+        vertices = (vertices - center) * scale
+        assert torch.all(vertices >= -0.5) and torch.all(vertices <= 0.5), 'vertices out of range'
+        return {'mesh': [Mesh(vertices=vertices, faces=faces)]}
+    
+    def read_dual_grid(self, root, instance):
+        coords, attr = o_voxel.io.read_vxz(os.path.join(root, f'{instance}.vxz'), num_threads=4)
+        vertices = sp.SparseTensor(
+            (attr['vertices'] / 255.0).float(),
+            torch.cat([torch.zeros_like(coords[:, 0:1]), coords], dim=-1),
+        )
+        intersected = vertices.replace(torch.cat([
+            attr['intersected'] % 2,
+            attr['intersected'] // 2 % 2,
+            attr['intersected'] // 4 % 2,
+        ], dim=-1).bool())
+        return {'vertices': vertices, 'intersected': intersected}
+
+    def get_instance(self, root, instance):
+        mesh = self.read_mesh(root['mesh_dump'], instance)
+        dual_grid = self.read_dual_grid(root['dual_grid'], instance)
+        return {**mesh, **dual_grid}
+    
+    @staticmethod
+    def collate_fn(batch, split_size=None):
+        if split_size is None:
+            group_idx = [list(range(len(batch)))]
+        else:
+            group_idx = load_balanced_group_indices([b['vertices'].feats.shape[0] for b in batch], split_size)
+        packs = []
+        for group in group_idx:
+            sub_batch = [batch[i] for i in group]
+            pack = {}
+
+            keys = [k for k in sub_batch[0].keys()]
+            for k in keys:
+                if isinstance(sub_batch[0][k], torch.Tensor):
+                    pack[k] = torch.stack([b[k] for b in sub_batch])
+                elif isinstance(sub_batch[0][k], sp.SparseTensor):
+                    pack[k] = sp.sparse_cat([b[k] for b in sub_batch], dim=0)
+                elif isinstance(sub_batch[0][k], list):
+                    pack[k] = sum([b[k] for b in sub_batch], [])
+                else:
+                    pack[k] = [b[k] for b in sub_batch]
+            
+            packs.append(pack)
+        
+        if split_size is None:
+            return packs[0]
+        return packs
+    
\ No newline at end of file
diff --git a/pixal3d/datasets/sparse_structure_latent.py b/pixal3d/datasets/sparse_structure_latent.py
new file mode 100644
index 0000000..c076067
--- /dev/null
+++ b/pixal3d/datasets/sparse_structure_latent.py
@@ -0,0 +1,408 @@
+import os
+import json
+from typing import *
+import numpy as np
+import torch
+import utils3d
+from PIL import Image
+from ..representations import Voxel
+from ..renderers import VoxelRenderer
+from .components import StandardDatasetBase, ImageConditionedMixin, ViewImageConditionedMixin
+from .. import models
+from ..utils.render_utils import yaw_pitch_r_fov_to_extrinsics_intrinsics
+
+
+class SparseStructureLatentVisMixin:
+    def __init__(
+        self,
+        *args,
+        pretrained_ss_dec: str = 'JeffreyXiang/TRELLIS-image-large/ckpts/ss_dec_conv3d_16l8_fp16.json',
+        ss_dec_path: Optional[str] = None,
+        ss_dec_ckpt: Optional[str] = None,
+        **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        self.ss_dec = None
+        self.pretrained_ss_dec = pretrained_ss_dec
+        self.ss_dec_path = ss_dec_path
+        self.ss_dec_ckpt = ss_dec_ckpt
+        
+    def _loading_ss_dec(self):
+        if self.ss_dec is not None:
+            return
+        if self.ss_dec_path is not None:
+            cfg = json.load(open(os.path.join(self.ss_dec_path, 'config.json'), 'r'))
+            decoder = getattr(models, cfg['models']['decoder']['name'])(**cfg['models']['decoder']['args'])
+            ckpt_path = os.path.join(self.ss_dec_path, 'ckpts', f'decoder_{self.ss_dec_ckpt}.pt')
+            decoder.load_state_dict(torch.load(ckpt_path, map_location='cpu', weights_only=True))
+        else:
+            decoder = models.from_pretrained(self.pretrained_ss_dec)
+        self.ss_dec = decoder.cuda().eval()
+
+    def _delete_ss_dec(self):
+        del self.ss_dec
+        self.ss_dec = None
+
+    @torch.no_grad()
+    def decode_latent(self, z, batch_size=4):
+        self._loading_ss_dec()
+        ss = []
+        if self.normalization:
+            z = z * self.std.to(z.device) + self.mean.to(z.device)
+        for i in range(0, z.shape[0], batch_size):
+            ss.append(self.ss_dec(z[i:i+batch_size]))
+        ss = torch.cat(ss, dim=0)
+        self._delete_ss_dec()
+        return ss
+
+    @torch.no_grad()
+    def visualize_sample(
+        self, 
+        x_0: Union[torch.Tensor, dict],
+        camera_angle_x: Optional[torch.Tensor] = None,
+        camera_distance: Optional[torch.Tensor] = None,
+        mesh_scale: Optional[torch.Tensor] = None,
+    ):
+        """
+        Visualize sparse structure samples.
+        
+        Args:
+            x_0: Latent tensor [B, C, D, H, W] or dict containing 'x_0'
+            camera_angle_x: Optional [B] camera FOV angle in radians
+            camera_distance: Optional [B] camera distance for GT view rendering
+            mesh_scale: Optional [B] mesh scale factor for coordinate alignment
+            
+        Returns:
+            dict with:
+                'multiview': [B, 3, 1024, 1024] - 4 fixed views rendered in 2x2 grid
+                'gt_view': [B, 3, 512, 512] - GT camera view (if camera params provided)
+        """
+        x_0 = x_0 if isinstance(x_0, torch.Tensor) else x_0['x_0']
+        x_0 = self.decode_latent(x_0.cuda())
+        
+        renderer = VoxelRenderer()
+        renderer.rendering_options.resolution = 512
+        renderer.rendering_options.ssaa = 4
+        
+        # Build fixed camera views (4 views: 0°, 90°, 180°, 270°)
+        yaw = [0, np.pi/2, np.pi, 3*np.pi/2]
+        yaw_offset = -16 / 180 * np.pi
+        yaw = [y + yaw_offset for y in yaw]
+        pitch = [20 / 180 * np.pi for _ in range(4)]
+        fixed_exts, fixed_ints = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaw, pitch, 2, 30)
+
+        # Check if we have GT camera parameters for front view rendering
+        # GT view uses the fixed front_view_transform_matrix from image_conditioned_proj.py
+        has_gt_camera = (
+            camera_angle_x is not None and 
+            camera_distance is not None and 
+            mesh_scale is not None
+        )
+        
+        multiview_images = []
+        gt_view_images = []
+        
+        # Build each representation
+        x_0 = x_0.cuda()
+        for i in range(x_0.shape[0]):
+            coords = torch.nonzero(x_0[i, 0] > 0, as_tuple=False)
+            resolution = x_0.shape[-1]
+            color = coords / resolution
+            
+            # Standard voxel for fixed multiview rendering (origin at [-0.5, -0.5, -0.5])
+            rep = Voxel(
+                origin=[-0.5, -0.5, -0.5],
+                voxel_size=1/resolution,
+                coords=coords,
+                attrs=color,
+                layout={
+                    'color': slice(0, 3),
+                }
+            )
+            
+            # Render 4 fixed views (2x2 grid)
+            image = torch.zeros(3, 1024, 1024).cuda()
+            tile = [2, 2]
+            for j, (ext, intr) in enumerate(zip(fixed_exts, fixed_ints)):
+                res = renderer.render(rep, ext, intr, colors_overwrite=color)
+                image[:, 512 * (j // tile[1]):512 * (j // tile[1] + 1), 512 * (j % tile[1]):512 * (j % tile[1] + 1)] = res['color']
+            multiview_images.append(image)
+            
+            # Render GT camera view using the fixed front view from image_conditioned_proj.py
+            if has_gt_camera:
+                # The GT view should match exactly how ProjGrid projects 3D points to 2D.
+                # 
+                # In image_conditioned_proj.py (ProjGrid.forward):
+                # 1. grid_points are in [-1, 1]^3 (from torch.linspace(-1, 1, res))
+                # 2. grid_points are rotated by rotation_matrix (Y-Z swap): x'=x, y'=-z, z'=y
+                # 3. grid_points are scaled: grid_points / mesh_scale / 2
+                # 4. Points are projected using front_view_transform_matrix with distance
+                #
+                # front_view_transform_matrix (camera-to-world):
+                # [[1, 0, 0, 0],
+                #  [0, 0, -1, -distance],
+                #  [0, 1, 0, 0],
+                #  [0, 0, 0, 1]]
+                # 
+                # Camera is at (0, -distance, 0) in Blender coords (Z-up), looking at origin.
+                #
+                # To match this in VoxelRenderer:
+                # 1. Voxel coords [0, res-1] map to positions via: pos = (coords + 0.5) * voxel_size + origin
+                # 2. We need these positions to match ProjGrid's transformed grid_points
+                # 3. Apply rotation by swapping/flipping coords, then scale voxel_size and origin
+                
+                scale = mesh_scale[i].item()
+                distance = camera_distance[i].item()
+                fov = camera_angle_x[i].item()
+                
+                # Coordinate transformation to match ProjGrid's rotation (x'=x, y'=-z, z'=y)
+                # new_coords maps to rotated positions in the same grid structure
+                new_coords = torch.zeros_like(coords)
+                new_coords[:, 0] = coords[:, 0]                       # x stays
+                new_coords[:, 1] = (resolution - 1) - coords[:, 2]    # y' = -z (flip for negation)
+                new_coords[:, 2] = coords[:, 1]                       # z' = y
+                
+                # Voxel position calculation:
+                # Original: pos = (coords + 0.5) / res - 0.5  -> range [-0.5, 0.5]
+                # We need:  pos = (coords + 0.5) * 2 / res - 1 -> range [-1, 1] (like ProjGrid)
+                # Then:     pos_final = pos / scale / 2        -> range [-0.5/scale, 0.5/scale]
+                #
+                # Combined: pos_final = ((coords + 0.5) * 2 / res - 1) / scale / 2
+                #                     = (coords + 0.5) / res / scale - 0.5 / scale
+                #                     = (coords + 0.5) * voxel_size + origin
+                # where:    voxel_size = 1 / res / scale
+                #           origin = -0.5 / scale
+                
+                scaled_voxel_size = 1.0 / resolution / scale
+                scaled_origin = [-0.5 / scale, -0.5 / scale, -0.5 / scale]
+                
+                rep_scaled = Voxel(
+                    origin=scaled_origin,
+                    voxel_size=scaled_voxel_size,
+                    coords=new_coords,
+                    attrs=color,
+                    layout={
+                        'color': slice(0, 3),
+                    }
+                )
+                
+                # Build the fixed front view camera (same as front_view_transform_matrix)
+                # Camera at (0, -distance, 0), looking at origin, up is Z
+                cam_pos = torch.tensor([0.0, -distance, 0.0], device=coords.device)
+                look_at = torch.tensor([0.0, 0.0, 0.0], device=coords.device)
+                cam_up = torch.tensor([0.0, 0.0, 1.0], device=coords.device)
+                
+                gt_ext = utils3d.torch.extrinsics_look_at(cam_pos, look_at, cam_up)
+                gt_int = utils3d.torch.intrinsics_from_fov_xy(
+                    torch.tensor(fov, device=coords.device),
+                    torch.tensor(fov, device=coords.device)
+                )
+                
+                # Ensure tensors are on the correct device (utils3d may not preserve device)
+                gt_ext = gt_ext.to(coords.device)
+                gt_int = gt_int.to(coords.device)
+                
+                gt_res = renderer.render(rep_scaled, gt_ext, gt_int, colors_overwrite=color)
+                gt_view_images.append(gt_res['color'])
+        
+        result = {
+            'multiview': torch.stack(multiview_images),
+        }
+        
+        if has_gt_camera and len(gt_view_images) > 0:
+            result['gt_view'] = torch.stack(gt_view_images)
+            
+        return result
+
+
+class SparseStructureLatent(SparseStructureLatentVisMixin, StandardDatasetBase):
+    """
+    Sparse structure latent dataset
+    
+    Args:
+        roots (str): path to the dataset
+        min_aesthetic_score (float): minimum aesthetic score
+        normalization (dict): normalization stats
+        pretrained_ss_dec (str): name of the pretrained sparse structure decoder
+        ss_dec_path (str): path to the sparse structure decoder, if given, will override the pretrained_ss_dec
+        ss_dec_ckpt (str): name of the sparse structure decoder checkpoint
+        skip_list (str, optional): path to a file containing sha256 hashes to skip
+        skip_aesthetic_score_datasets (list, optional): list of dataset names to skip aesthetic score check
+    """
+    def __init__(self,
+        roots: str,
+        *,
+        min_aesthetic_score: float = 5.0,
+        normalization: Optional[dict] = None,
+        pretrained_ss_dec: str = 'JeffreyXiang/TRELLIS-image-large/ckpts/ss_dec_conv3d_16l8_fp16',
+        ss_dec_path: Optional[str] = None,
+        ss_dec_ckpt: Optional[str] = None,
+        skip_list: Optional[str] = None,
+        skip_aesthetic_score_datasets: Optional[list] = None,
+    ):
+        self.min_aesthetic_score = min_aesthetic_score
+        self.normalization = normalization
+        self.value_range = (0, 1)
+        
+        super().__init__(
+            roots,
+            pretrained_ss_dec=pretrained_ss_dec,
+            ss_dec_path=ss_dec_path,
+            ss_dec_ckpt=ss_dec_ckpt,
+            skip_list=skip_list,
+            skip_aesthetic_score_datasets=skip_aesthetic_score_datasets,
+        )
+        
+        if self.normalization is not None:
+            self.mean = torch.tensor(self.normalization['mean']).reshape(-1, 1, 1, 1)
+            self.std = torch.tensor(self.normalization['std']).reshape(-1, 1, 1, 1)
+  
+    def filter_metadata(self, metadata, dataset_name=None):
+        stats = {}
+        metadata = metadata[metadata['ss_latent_encoded'] == True]
+        stats['With latent'] = len(metadata)
+        # Skip aesthetic score check for specified datasets (e.g., texverse) or if column doesn't exist
+        skip_aesthetic = (
+            (dataset_name and dataset_name.lower() in [d.lower() for d in self.skip_aesthetic_score_datasets]) or
+            ('aesthetic_score' not in metadata.columns)
+        )
+        if skip_aesthetic:
+            stats[f'Aesthetic score check skipped'] = len(metadata)
+        else:
+            metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]
+            stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)
+        return metadata, stats
+                
+    def get_instance(self, root, instance):
+        latent = np.load(os.path.join(root['ss_latent'], f'{instance}.npz'))
+        z = torch.tensor(latent['z']).float()
+        if self.normalization is not None:
+            z = (z - self.mean) / self.std
+
+        pack = {
+            'x_0': z,
+        }
+        return pack
+
+
+class ImageConditionedSparseStructureLatent(ImageConditionedMixin, SparseStructureLatent):
+    """
+    Image-conditioned sparse structure dataset
+    """
+    pass
+
+
+class SparseStructureLatentView(SparseStructureLatentVisMixin, StandardDatasetBase):
+    """
+    View-based sparse structure latent dataset.
+    
+    Data format: {sha256}/view{XX}.npz where each npz contains 'z' key.
+    
+    Args:
+        num_views (int): Number of views to use (0 to num_views-1). Default is 2.
+        skip_list (str, optional): path to a file containing sha256 hashes to skip
+        skip_aesthetic_score_datasets (list, optional): list of dataset names to skip aesthetic score check
+    """
+    def __init__(self,
+        roots: str,
+        *,
+        min_aesthetic_score: float = 5.0,
+        normalization: Optional[dict] = None,
+        num_views: int = 2,
+        pretrained_ss_dec: str = 'JeffreyXiang/TRELLIS-image-large/ckpts/ss_dec_conv3d_16l8_fp16',
+        ss_dec_path: Optional[str] = None,
+        ss_dec_ckpt: Optional[str] = None,
+        skip_list: Optional[str] = None,
+        skip_aesthetic_score_datasets: Optional[list] = None,
+    ):
+        self.min_aesthetic_score = min_aesthetic_score
+        self.normalization = normalization
+        self.num_views = num_views
+        self.value_range = (0, 1)
+        
+        super().__init__(
+            roots,
+            pretrained_ss_dec=pretrained_ss_dec,
+            ss_dec_path=ss_dec_path,
+            ss_dec_ckpt=ss_dec_ckpt,
+            skip_list=skip_list,
+            skip_aesthetic_score_datasets=skip_aesthetic_score_datasets,
+        )
+        
+        if self.normalization is not None:
+            self.mean = torch.tensor(self.normalization['mean']).reshape(-1, 1, 1, 1)
+            self.std = torch.tensor(self.normalization['std']).reshape(-1, 1, 1, 1)
+  
+    def filter_metadata(self, metadata, dataset_name=None):
+        stats = {}
+        # View-based ss_latent uses columns like:
+        #   ss_latent_view00_encoded, ss_latent_view01_encoded, ... (view format)
+        #   ss_latent_view_scale00_encoded, ss_latent_view_scale01_encoded, ... (view_scale format)
+        # Check both formats and use whichever exists (prefer view_scale over view)
+        required_view_cols = [f'ss_latent_view_scale{i:02d}_encoded' for i in range(self.num_views)]
+        existing_view_cols = [col for col in required_view_cols if col in metadata.columns]
+        
+        if not existing_view_cols:
+            # Fallback to view format
+            required_view_cols = [f'ss_latent_view{i:02d}_encoded' for i in range(self.num_views)]
+            existing_view_cols = [col for col in required_view_cols if col in metadata.columns]
+        
+        if existing_view_cols:
+            # Filter rows where all required views are encoded
+            # Note: NaN should be treated as False, so use == True for explicit comparison
+            has_all_views = (metadata[existing_view_cols] == True).all(axis=1)
+            metadata = metadata[has_all_views]
+            stats[f'With {self.num_views} view latents'] = len(metadata)
+        else:
+            # Fallback: check ss_latent_encoded column
+            if 'ss_latent_encoded' in metadata.columns:
+                metadata = metadata[metadata['ss_latent_encoded'] == True]
+                stats['With latent'] = len(metadata)
+            else:
+                raise ValueError(f'No view columns found in metadata: {metadata.columns.tolist()}')
+        # Skip aesthetic score check for specified datasets (e.g., texverse) or if column doesn't exist
+        skip_aesthetic = (
+            (dataset_name and dataset_name.lower() in [d.lower() for d in self.skip_aesthetic_score_datasets]) or
+            ('aesthetic_score' not in metadata.columns)
+        )
+        if skip_aesthetic:
+            stats[f'Aesthetic score check skipped'] = len(metadata)
+        else:
+            metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]
+            stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)
+        return metadata, stats
+                
+    def get_instance(self, root, instance):
+        # View-based format: directory with view{XX}.npz files
+        latent_dir = os.path.join(root['ss_latent'], instance)
+        
+        # Randomly select a view from the configured range
+        view_idx = np.random.randint(0, self.num_views)
+        view_file = f'view{view_idx:02d}.npz'
+        
+        # Store view info for ViewImageConditionedMixin
+        self._current_view_idx = view_idx
+        self._current_latent_dir = latent_dir
+        
+        latent = np.load(os.path.join(latent_dir, view_file))
+        z = torch.tensor(latent['z']).float()
+        if self.normalization is not None:
+            z = (z - self.mean) / self.std
+
+        pack = {
+            'x_0': z,
+            'view_idx': view_idx,
+        }
+        return pack
+
+
+class ViewImageConditionedSparseStructureLatentView(ViewImageConditionedMixin, SparseStructureLatentView):
+    """
+    Image-conditioned view-based sparse structure dataset.
+    
+    Loads ss_latent from {sha256}/view{XX}.npz format and pairs with 
+    corresponding view from render_cond.
+    
+    Uses ViewImageConditionedMixin which reads mesh_scale from view{XX}_scale.json.
+    """
+    pass
diff --git a/pixal3d/datasets/sparse_voxel_pbr.py b/pixal3d/datasets/sparse_voxel_pbr.py
new file mode 100644
index 0000000..d9146b1
--- /dev/null
+++ b/pixal3d/datasets/sparse_voxel_pbr.py
@@ -0,0 +1,298 @@
+import os
+import io
+from typing import Union
+import numpy as np
+import pickle
+import torch
+from PIL import Image
+import o_voxel
+import utils3d
+from .components import StandardDatasetBase
+from ..modules import sparse as sp
+from ..renderers import VoxelRenderer
+from ..representations import Voxel
+from ..representations.mesh import MeshWithPbrMaterial, TextureFilterMode, TextureWrapMode, AlphaMode, PbrMaterial, Texture
+
+from ..utils.data_utils import load_balanced_group_indices
+
+
+def is_power_of_two(n: int) -> bool:
+    return n > 0 and (n & (n - 1)) == 0
+
+
+def nearest_power_of_two(n: int) -> int:
+    if n < 1:
+        raise ValueError("n must be >= 1")
+    if is_power_of_two(n):
+        return n
+    lower = 2 ** (n.bit_length() - 1)
+    upper = 2 ** n.bit_length()
+    if n - lower < upper - n:
+        return lower
+    else:
+        return upper
+    
+
+class SparseVoxelPbrVisMixin:
+    @torch.no_grad()
+    def visualize_sample(self, x: Union[sp.SparseTensor, dict]):
+        x = x if isinstance(x, sp.SparseTensor) else x['x']
+        
+        renderer = VoxelRenderer()
+        renderer.rendering_options.resolution = 512
+        renderer.rendering_options.ssaa = 4
+        
+        # Build camera
+        yaws = [0, np.pi / 2, np.pi, 3 * np.pi / 2]
+        yaws_offset = np.random.uniform(-np.pi / 4, np.pi / 4)
+        yaws = [y + yaws_offset for y in yaws]
+        pitch = [np.random.uniform(-np.pi / 4, np.pi / 4) for _ in range(4)]
+
+        exts = []
+        ints = []
+        for yaw, pitch in zip(yaws, pitch):
+            orig = torch.tensor([
+                np.sin(yaw) * np.cos(pitch),
+                np.cos(yaw) * np.cos(pitch),
+                np.sin(pitch),
+            ]).float().cuda() * 2
+            fov = torch.deg2rad(torch.tensor(30)).cuda()
+            extrinsics = utils3d.torch.extrinsics_look_at(orig, torch.tensor([0, 0, 0]).float().cuda(), torch.tensor([0, 0, 1]).float().cuda())
+            intrinsics = utils3d.torch.intrinsics_from_fov_xy(fov, fov)
+            exts.append(extrinsics)
+            ints.append(intrinsics)
+
+        images = {k: [] for k in self.layout}
+        
+        # Build each representation
+        x = x.cuda()
+        for i in range(x.shape[0]):
+            rep = Voxel(
+                origin=[-0.5, -0.5, -0.5],
+                voxel_size=1/self.resolution,
+                coords=x[i].coords[:, 1:].contiguous(),
+                attrs=None,
+                layout={
+                    'color': slice(0, 3),
+                }
+            )
+            for k in self.layout:
+                image = torch.zeros(3, 1024, 1024).cuda()
+                tile = [2, 2]
+                for j, (ext, intr) in enumerate(zip(exts, ints)):
+                    attr = x[i].feats[:, self.layout[k]].expand(-1, 3)
+                    res = renderer.render(rep, ext, intr, colors_overwrite=attr)
+                    image[:, 512 * (j // tile[1]):512 * (j // tile[1] + 1), 512 * (j % tile[1]):512 * (j % tile[1] + 1)] = res['color']
+                images[k].append(image)
+        
+        for k in self.layout:
+            images[k] = torch.stack(images[k])
+            
+        return images
+
+
+class SparseVoxelPbrDataset(SparseVoxelPbrVisMixin, StandardDatasetBase):
+    """
+    Sparse Voxel PBR dataset.
+    
+    Args:
+        roots (str): path to the dataset
+        resolution (int): resolution of the voxel grid
+        min_aesthetic_score (float): minimum aesthetic score of the instances to be included in the dataset
+    """
+
+    def __init__(
+        self,
+        roots,
+        resolution: int = 1024,
+        max_active_voxels: int = 1000000,
+        max_num_faces: int = None,
+        min_aesthetic_score: float = 5.0,
+        attrs: list[str] = ['base_color', 'metallic', 'roughness', 'emissive', 'alpha'],
+        with_mesh: bool = True,
+    ):
+        self.resolution = resolution
+        self.min_aesthetic_score = min_aesthetic_score
+        self.max_active_voxels = max_active_voxels
+        self.max_num_faces = max_num_faces
+        self.with_mesh = with_mesh
+        self.value_range = (-1, 1)
+        self.channels = {
+            'base_color': 3,
+            'metallic': 1,
+            'roughness': 1,
+            'emissive': 3,
+            'alpha': 1,
+        }
+        self.layout = {}
+        start = 0
+        for attr in attrs:
+            self.layout[attr] = slice(start, start + self.channels[attr])
+            start += self.channels[attr]
+
+        super().__init__(roots)
+        
+        self.loads = [self.metadata.loc[sha256, f'num_pbr_voxels'] for _, sha256, _ in self.instances]
+        
+    def __str__(self):
+        lines = [
+            super().__str__(),
+            f'  - Resolution: {self.resolution}',
+            f'  - Attributes: {list(self.layout.keys())}',
+        ]
+        return '\n'.join(lines)
+        
+    def filter_metadata(self, metadata, dataset_name=None):
+        stats = {}
+        metadata = metadata[metadata['pbr_voxelized'] == True]
+        stats['PBR Voxelized'] = len(metadata)
+        if self.min_aesthetic_score is not None:
+            metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]
+            stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)
+        metadata = metadata[metadata['num_pbr_voxels'] <= self.max_active_voxels]
+        stats[f'Active voxels <= {self.max_active_voxels}'] = len(metadata)
+        if self.max_num_faces is not None:
+            metadata = metadata[metadata['num_faces'] <= self.max_num_faces]
+            stats[f'Faces <= {self.max_num_faces}'] = len(metadata)
+        return metadata, stats
+
+    @staticmethod
+    def _texture_from_dump(pack) -> Texture:
+        png_bytes = pack['image']
+        image = Image.open(io.BytesIO(png_bytes))
+        if image.width != image.height or not is_power_of_two(image.width):
+            size = nearest_power_of_two(max(image.width, image.height))
+            image = image.resize((size, size), Image.LANCZOS)
+        texture = torch.tensor(np.array(image) / 255.0, dtype=torch.float32).reshape(image.height, image.width, -1)
+        filter_mode = {
+            'Linear': TextureFilterMode.LINEAR,
+            'Closest': TextureFilterMode.CLOSEST,
+            'Cubic': TextureFilterMode.LINEAR,
+            'Smart': TextureFilterMode.LINEAR,
+        }[pack['interpolation']]
+        wrap_mode = {
+            'REPEAT': TextureWrapMode.REPEAT,
+            'EXTEND': TextureWrapMode.CLAMP_TO_EDGE,
+            'CLIP': TextureWrapMode.CLAMP_TO_EDGE,
+            'MIRROR': TextureWrapMode.MIRRORED_REPEAT,
+        }[pack['extension']]
+        return Texture(texture, filter_mode=filter_mode, wrap_mode=wrap_mode)
+
+    def read_mesh_with_texture(self, root, instance):
+        with open(os.path.join(root, f'{instance}.pickle'), 'rb') as f:
+            dump = pickle.load(f)
+            
+        # Fix dump alpha map
+        for mat in dump['materials']:
+            if mat['alphaTexture'] is not None and mat['alphaMode'] == 'OPAQUE':
+                mat['alphaMode'] = 'BLEND'
+
+        # process material
+        materials = []
+        for mat in dump['materials']:
+            materials.append(PbrMaterial(
+                base_color_texture=self._texture_from_dump(mat['baseColorTexture']) if mat['baseColorTexture'] is not None else None,
+                base_color_factor=mat['baseColorFactor'],
+                metallic_texture=self._texture_from_dump(mat['metallicTexture']) if mat['metallicTexture'] is not None else None,
+                metallic_factor=mat['metallicFactor'],
+                roughness_texture=self._texture_from_dump(mat['roughnessTexture']) if mat['roughnessTexture'] is not None else None,
+                roughness_factor=mat['roughnessFactor'],
+                alpha_texture=self._texture_from_dump(mat['alphaTexture']) if mat['alphaTexture'] is not None else None,
+                alpha_factor=mat['alphaFactor'],
+                alpha_mode={
+                    'OPAQUE': AlphaMode.OPAQUE,
+                    'MASK': AlphaMode.MASK,
+                    'BLEND': AlphaMode.BLEND,
+                }[mat['alphaMode']],
+                alpha_cutoff=mat['alphaCutoff'],
+            ))
+        materials.append(PbrMaterial(
+            base_color_factor=[0.8, 0.8, 0.8],
+            alpha_factor=1.0,
+            metallic_factor=0.0,
+            roughness_factor=0.5,
+            alpha_mode=AlphaMode.OPAQUE,
+            alpha_cutoff=0.5,
+        ))  # append default material
+
+        # process mesh
+        start = 0
+        vertices = []
+        faces = []
+        material_ids = []
+        uv_coords = []
+        for obj in dump['objects']:
+            if obj['vertices'].size == 0 or obj['faces'].size == 0:
+                continue
+            vertices.append(obj['vertices'])
+            faces.append(obj['faces'] + start)
+            obj['mat_ids'][obj['mat_ids'] == -1] = len(materials) - 1
+            material_ids.append(obj['mat_ids'])
+            uv_coords.append(obj['uvs'] if obj['uvs'] is not None else np.zeros((obj['faces'].shape[0], 3, 2), dtype=np.float32))
+            start += len(obj['vertices'])
+        
+        vertices = torch.from_numpy(np.concatenate(vertices, axis=0)).float()
+        faces = torch.from_numpy(np.concatenate(faces, axis=0)).long()
+        material_ids = torch.from_numpy(np.concatenate(material_ids, axis=0)).long()
+        uv_coords = torch.from_numpy(np.concatenate(uv_coords, axis=0)).float()
+        
+        # Normalize vertices
+        vertices_min = vertices.min(dim=0)[0]
+        vertices_max = vertices.max(dim=0)[0]
+        center = (vertices_min + vertices_max) / 2
+        scale = 0.99999 / (vertices_max - vertices_min).max()
+        vertices = (vertices - center) * scale
+        assert torch.all(vertices >= -0.5) and torch.all(vertices <= 0.5), 'vertices out of range'
+        
+        return {'mesh': [MeshWithPbrMaterial(
+            vertices=vertices,
+            faces=faces,
+            material_ids=material_ids,
+            uv_coords=uv_coords,
+            materials=materials,
+        )]}
+
+    def read_pbr_voxel(self, root, instance):
+        coords, attr = o_voxel.io.read_vxz(os.path.join(root, f'{instance}.vxz'), num_threads=4)
+        feats = torch.concat([attr[k] for k in self.layout], dim=-1) / 255.0 * 2 - 1
+        x = sp.SparseTensor(
+            feats.float(),
+            torch.cat([torch.zeros_like(coords[:, 0:1]), coords], dim=-1),
+        )
+        return {'x': x}
+    
+    def get_instance(self, root, instance):
+        if self.with_mesh:
+            mesh = self.read_mesh_with_texture(root['pbr_dump'], instance)
+            pbr_voxel = self.read_pbr_voxel(root['pbr_voxel'], instance)
+            return {**mesh, **pbr_voxel}
+        else:
+            return self.read_pbr_voxel(root['pbr_voxel'], instance)
+    
+    @staticmethod
+    def collate_fn(batch, split_size=None):
+        if split_size is None:
+            group_idx = [list(range(len(batch)))]
+        else:
+            group_idx = load_balanced_group_indices([b['x'].feats.shape[0] for b in batch], split_size)
+        packs = []
+        for group in group_idx:
+            sub_batch = [batch[i] for i in group]
+            pack = {}
+
+            keys = [k for k in sub_batch[0].keys()]
+            for k in keys:
+                if isinstance(sub_batch[0][k], torch.Tensor):
+                    pack[k] = torch.stack([b[k] for b in sub_batch])
+                elif isinstance(sub_batch[0][k], sp.SparseTensor):
+                    pack[k] = sp.sparse_cat([b[k] for b in sub_batch], dim=0)
+                elif isinstance(sub_batch[0][k], list):
+                    pack[k] = sum([b[k] for b in sub_batch], [])
+                else:
+                    pack[k] = [b[k] for b in sub_batch]
+            
+            packs.append(pack)
+        
+        if split_size is None:
+            return packs[0]
+        return packs
diff --git a/pixal3d/datasets/structured_latent.py b/pixal3d/datasets/structured_latent.py
new file mode 100644
index 0000000..8cc0bcc
--- /dev/null
+++ b/pixal3d/datasets/structured_latent.py
@@ -0,0 +1,224 @@
+import json
+import os
+from typing import *
+import numpy as np
+import torch
+import utils3d.torch
+from .components import StandardDatasetBase, ImageConditionedMixin
+from ..modules.sparse.basic import SparseTensor
+from .. import models
+from ..utils.render_utils import get_renderer
+from ..utils.data_utils import load_balanced_group_indices
+
+
+class SLatVisMixin:
+    def __init__(
+        self,
+        *args,
+        pretrained_slat_dec: str = 'JeffreyXiang/TRELLIS-image-large/ckpts/slat_dec_gs_swin8_B_64l8gs32_fp16',
+        slat_dec_path: Optional[str] = None,
+        slat_dec_ckpt: Optional[str] = None,
+        **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        self.slat_dec = None
+        self.pretrained_slat_dec = pretrained_slat_dec
+        self.slat_dec_path = slat_dec_path
+        self.slat_dec_ckpt = slat_dec_ckpt
+        
+    def _loading_slat_dec(self):
+        if self.slat_dec is not None:
+            return
+        if self.slat_dec_path is not None:
+            cfg = json.load(open(os.path.join(self.slat_dec_path, 'config.json'), 'r'))
+            decoder = getattr(models, cfg['models']['decoder']['name'])(**cfg['models']['decoder']['args'])
+            ckpt_path = os.path.join(self.slat_dec_path, 'ckpts', f'decoder_{self.slat_dec_ckpt}.pt')
+            decoder.load_state_dict(torch.load(ckpt_path, map_location='cpu', weights_only=True))
+        else:
+            decoder = models.from_pretrained(self.pretrained_slat_dec)
+        self.slat_dec = decoder.cuda().eval()
+
+    def _delete_slat_dec(self):
+        del self.slat_dec
+        self.slat_dec = None
+
+    @torch.no_grad()
+    def decode_latent(self, z, batch_size=4):
+        self._loading_slat_dec()
+        reps = []
+        if self.normalization is not None:
+            z = z * self.std.to(z.device) + self.mean.to(z.device)
+        for i in range(0, z.shape[0], batch_size):
+            reps.append(self.slat_dec(z[i:i+batch_size]))
+        reps = sum(reps, [])
+        self._delete_slat_dec()
+        return reps
+
+    @torch.no_grad()
+    def visualize_sample(self, x_0: Union[SparseTensor, dict]):
+        x_0 = x_0 if isinstance(x_0, SparseTensor) else x_0['x_0']
+        reps = self.decode_latent(x_0.cuda())
+        
+        # Build camera
+        yaws = [0, np.pi / 2, np.pi, 3 * np.pi / 2]
+        yaws_offset = np.random.uniform(-np.pi / 4, np.pi / 4)
+        yaws = [y + yaws_offset for y in yaws]
+        pitch = [np.random.uniform(-np.pi / 4, np.pi / 4) for _ in range(4)]
+
+        exts = []
+        ints = []
+        for yaw, pitch in zip(yaws, pitch):
+            orig = torch.tensor([
+                np.sin(yaw) * np.cos(pitch),
+                np.cos(yaw) * np.cos(pitch),
+                np.sin(pitch),
+            ]).float().cuda() * 2
+            fov = torch.deg2rad(torch.tensor(40)).cuda()
+            extrinsics = utils3d.torch.extrinsics_look_at(orig, torch.tensor([0, 0, 0]).float().cuda(), torch.tensor([0, 0, 1]).float().cuda())
+            intrinsics = utils3d.torch.intrinsics_from_fov_xy(fov, fov)
+            exts.append(extrinsics)
+            ints.append(intrinsics)
+
+        renderer = get_renderer(reps[0])
+        images = []
+        for representation in reps:
+            image = torch.zeros(3, 1024, 1024).cuda()
+            tile = [2, 2]
+            for j, (ext, intr) in enumerate(zip(exts, ints)):
+                res = renderer.render(representation, ext, intr)
+                image[:, 512 * (j // tile[1]):512 * (j // tile[1] + 1), 512 * (j % tile[1]):512 * (j % tile[1] + 1)] = res['color']
+            images.append(image)
+        images = torch.stack(images)
+            
+        return images
+    
+    
+class SLat(SLatVisMixin, StandardDatasetBase):
+    """
+    structured latent V2 dataset
+    
+    Args:
+        roots (str): path to the dataset
+        min_aesthetic_score (float): minimum aesthetic score
+        max_tokens (int): maximum number of tokens
+        latent_key (str): key of the latent to be used
+        normalization (dict): normalization stats
+        pretrained_slat_dec (str): name of the pretrained slat decoder
+        slat_dec_path (str): path to the slat decoder, if given, will override the pretrained_slat_dec
+        slat_dec_ckpt (str): name of the slat decoder checkpoint
+        skip_list (str, optional): path to a file containing sha256 hashes to skip
+        skip_aesthetic_score_datasets (list, optional): list of dataset names to skip aesthetic score check
+    """
+    def __init__(self,
+        roots: str,
+        *,
+        min_aesthetic_score: float = 5.0,
+        max_tokens: int = 32768,
+        latent_key: str = 'shape_latent',
+        normalization: Optional[dict] = None,
+        pretrained_slat_dec: str = 'JeffreyXiang/TRELLIS-image-large/ckpts/slat_dec_gs_swin8_B_64l8gs32_fp16',
+        slat_dec_path: Optional[str] = None,
+        slat_dec_ckpt: Optional[str] = None,
+        skip_list: Optional[str] = None,
+        skip_aesthetic_score_datasets: Optional[list] = None,
+    ):
+        self.normalization = normalization
+        self.min_aesthetic_score = min_aesthetic_score
+        self.max_tokens = max_tokens
+        self.latent_key = latent_key
+        self.value_range = (0, 1)
+        
+        super().__init__(
+            roots,
+            pretrained_slat_dec=pretrained_slat_dec,
+            slat_dec_path=slat_dec_path,
+            slat_dec_ckpt=slat_dec_ckpt,
+            skip_list=skip_list,
+            skip_aesthetic_score_datasets=skip_aesthetic_score_datasets,
+        )
+
+        self.loads = [self.metadata.loc[sha256, f'{latent_key}_tokens'] for _, sha256, _ in self.instances]
+        
+        if self.normalization is not None:
+            self.mean = torch.tensor(self.normalization['mean']).reshape(1, -1)
+            self.std = torch.tensor(self.normalization['std']).reshape(1, -1)
+      
+    def filter_metadata(self, metadata, dataset_name=None):
+        stats = {}
+        metadata = metadata[metadata[f'{self.latent_key}_encoded'] == True]
+        stats['With latent'] = len(metadata)
+        # Skip aesthetic score check for specified datasets (e.g., texverse) or if column doesn't exist
+        skip_aesthetic = (
+            (dataset_name and dataset_name.lower() in [d.lower() for d in self.skip_aesthetic_score_datasets]) or
+            ('aesthetic_score' not in metadata.columns)
+        )
+        if skip_aesthetic:
+            stats[f'Aesthetic score check skipped'] = len(metadata)
+        else:
+            metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]
+            stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)
+        metadata = metadata[metadata[f'{self.latent_key}_tokens'] <= self.max_tokens]
+        stats[f'Num tokens <= {self.max_tokens}'] = len(metadata)
+        return metadata, stats
+
+    def get_instance(self, root, instance):
+        data = np.load(os.path.join(root[self.latent_key], f'{instance}.npz'))
+        coords = torch.tensor(data['coords']).int()
+        feats = torch.tensor(data['feats']).float()
+        if self.normalization is not None:
+            feats = (feats - self.mean) / self.std
+        return {
+            'coords': coords,
+            'feats': feats,
+        }
+        
+    @staticmethod
+    def collate_fn(batch, split_size=None):
+        if split_size is None:
+            group_idx = [list(range(len(batch)))]
+        else:
+            group_idx = load_balanced_group_indices([b['coords'].shape[0] for b in batch], split_size)
+        packs = []
+        for group in group_idx:
+            sub_batch = [batch[i] for i in group]
+            pack = {}
+            coords = []
+            feats = []
+            layout = []
+            start = 0
+            for i, b in enumerate(sub_batch):
+                coords.append(torch.cat([torch.full((b['coords'].shape[0], 1), i, dtype=torch.int32), b['coords']], dim=-1))
+                feats.append(b['feats'])
+                layout.append(slice(start, start + b['coords'].shape[0]))
+                start += b['coords'].shape[0]
+            coords = torch.cat(coords)
+            feats = torch.cat(feats)
+            pack['x_0'] = SparseTensor(
+                coords=coords,
+                feats=feats,
+            )
+            pack['x_0']._shape = torch.Size([len(group), *sub_batch[0]['feats'].shape[1:]])
+            pack['x_0'].register_spatial_cache('layout', layout)
+            
+            # collate other data
+            keys = [k for k in sub_batch[0].keys() if k not in ['coords', 'feats']]
+            for k in keys:
+                if isinstance(sub_batch[0][k], torch.Tensor):
+                    pack[k] = torch.stack([b[k] for b in sub_batch])
+                elif isinstance(sub_batch[0][k], list):
+                    pack[k] = sum([b[k] for b in sub_batch], [])
+                else:
+                    pack[k] = [b[k] for b in sub_batch]
+                    
+            packs.append(pack)
+          
+        if split_size is None:
+            return packs[0]
+        return packs
+
+
+class ImageConditionedSLat(ImageConditionedMixin, SLat):
+    """
+    Image conditioned structured latent dataset
+    """
+    pass
diff --git a/pixal3d/datasets/structured_latent_shape.py b/pixal3d/datasets/structured_latent_shape.py
new file mode 100644
index 0000000..28a5511
--- /dev/null
+++ b/pixal3d/datasets/structured_latent_shape.py
@@ -0,0 +1,402 @@
+import os
+import json
+from typing import *
+import numpy as np
+import torch
+import utils3d
+from .. import models
+from .components import ImageConditionedMixin, ViewImageConditionedMixin
+from ..modules.sparse import SparseTensor
+from .structured_latent import SLatVisMixin, SLat
+from ..utils.render_utils import get_renderer, yaw_pitch_r_fov_to_extrinsics_intrinsics
+from ..utils.data_utils import load_balanced_group_indices
+
+
+class SLatShapeVisMixin(SLatVisMixin):
+    def _loading_slat_dec(self):
+        if self.slat_dec is not None:
+            return
+        if self.slat_dec_path is not None:
+            cfg = json.load(open(os.path.join(self.slat_dec_path, 'config.json'), 'r'))
+            decoder = getattr(models, cfg['models']['decoder']['name'])(**cfg['models']['decoder']['args'])
+            ckpt_path = os.path.join(self.slat_dec_path, 'ckpts', f'decoder_{self.slat_dec_ckpt}.pt')
+            decoder.load_state_dict(torch.load(ckpt_path, map_location='cpu', weights_only=True))
+        else:
+            decoder = models.from_pretrained(self.pretrained_slat_dec)
+        decoder.set_resolution(self.resolution)
+        self.slat_dec = decoder.cuda().eval()
+
+    @torch.no_grad()
+    def visualize_sample(
+        self, 
+        x_0: Union[SparseTensor, dict],
+        camera_angle_x: Optional[torch.Tensor] = None,
+        camera_distance: Optional[torch.Tensor] = None,
+        mesh_scale: Optional[torch.Tensor] = None,
+    ):
+        """
+        Visualize shape samples.
+        
+        Args:
+            x_0: SparseTensor or dict containing 'x_0'
+            camera_angle_x: Optional [B] camera FOV angle in radians
+            camera_distance: Optional [B] camera distance for GT view rendering
+            mesh_scale: Optional [B] mesh scale factor for coordinate alignment
+            
+        Returns:
+            dict with:
+                'multiview': [B, 3, 1024, 1024] - 4 fixed views rendered in 2x2 grid (normal)
+                'gt_view': [B, 3, 512, 512] - GT camera view (if camera params provided)
+        """
+        x_0 = x_0 if isinstance(x_0, SparseTensor) else x_0['x_0']
+        reps = self.decode_latent(x_0.cuda())
+        
+        # build fixed camera views (4 views: 0°, 90°, 180°, 270°)
+        yaw = [0, np.pi/2, np.pi, 3*np.pi/2]
+        yaw_offset = -16 / 180 * np.pi
+        yaw = [y + yaw_offset for y in yaw]
+        pitch = [20 / 180 * np.pi for _ in range(4)]
+        fixed_exts, fixed_ints = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaw, pitch, 2, 30)
+        
+        # Check if we have GT camera parameters for GT view rendering
+        has_gt_camera = (
+            camera_angle_x is not None and 
+            camera_distance is not None and 
+            mesh_scale is not None
+        )
+        
+        # render
+        renderer = get_renderer(reps[0])
+        multiview_images = []
+        gt_view_images = []
+        
+        for i, representation in enumerate(reps):
+            # Render 4 fixed views (2x2 grid)
+            image = torch.zeros(3, 1024, 1024).cuda()
+            tile = [2, 2]
+            
+            # Validate mesh data before rasterization
+            verts = representation.vertices
+            faces = representation.faces
+            if verts.shape[0] == 0 or faces.shape[0] == 0:
+                print(f"[visualize_sample] Warning: sample {i} has empty mesh, skipping")
+                multiview_images.append(image)
+                continue
+            if faces.max() >= verts.shape[0]:
+                print(f"[visualize_sample] Warning: sample {i} has out-of-bound face indices "
+                      f"(max face idx={faces.max().item()}, num verts={verts.shape[0]}), skipping")
+                multiview_images.append(image)
+                continue
+            if torch.isnan(verts).any() or torch.isinf(verts).any():
+                print(f"[visualize_sample] Warning: sample {i} has NaN/Inf vertices, skipping")
+                multiview_images.append(image)
+                continue
+            
+            try:
+                for j, (ext, intr) in enumerate(zip(fixed_exts, fixed_ints)):
+                    res = renderer.render(representation, ext, intr)
+                    image[:, 512 * (j // tile[1]):512 * (j // tile[1] + 1), 512 * (j % tile[1]):512 * (j % tile[1] + 1)] = res['normal']
+            except RuntimeError as e:
+                print(f"[visualize_sample] Warning: render failed for sample {i}: {e}")
+                image = torch.zeros(3, 1024, 1024).cuda()
+            multiview_images.append(image)
+            
+            # Render GT camera view using the fixed front view (same as sparse_structure_latent.py)
+            if has_gt_camera:
+                # The GT view should match exactly how ProjGrid projects 3D points to 2D.
+                # 
+                # In image_conditioned_proj.py (ProjGrid.forward):
+                # 1. grid_points are in [-1, 1]^3 (from torch.linspace(-1, 1, res))
+                # 2. grid_points are rotated by rotation_matrix (Y-Z swap): x'=x, y'=-z, z'=y
+                # 3. grid_points are scaled: grid_points / mesh_scale / 2
+                # 4. Points are projected using front_view_transform_matrix with distance
+                #
+                # Mesh vertices are in [-0.5, 0.5]^3. To match ProjGrid's coordinate space,
+                # we need to scale them: vertices / mesh_scale -> [-0.5/s, 0.5/s]^3
+                # This is equivalent to ProjGrid's: [-1,1]^3 / scale / 2 -> [-0.5/s, 0.5/s]^3
+                #
+                # Camera position: ProjGrid camera is at (0, -distance, 0) in Blender coords (Z-up).
+                # After inverse rotation to mesh space, camera is at (0, 0, distance).
+                
+                scale = mesh_scale[i].item()
+                distance = camera_distance[i].item()
+                fov = camera_angle_x[i].item()
+                device = representation.vertices.device
+                
+                # Scale mesh vertices to match ProjGrid's projection space
+                from ..representations import Mesh
+                scaled_rep = Mesh(
+                    vertices=representation.vertices / scale,
+                    faces=representation.faces,
+                )
+                
+                cam_pos = torch.tensor([0.0, 0.0, distance], device=device)
+                look_at = torch.tensor([0.0, 0.0, 0.0], device=device)
+                cam_up = torch.tensor([0.0, 1.0, 0.0], device=device)
+                
+                gt_ext = utils3d.torch.extrinsics_look_at(cam_pos, look_at, cam_up)
+                gt_int = utils3d.torch.intrinsics_from_fov_xy(
+                    torch.tensor(fov, device=device),
+                    torch.tensor(fov, device=device)
+                )
+                
+                gt_ext = gt_ext.to(device)
+                gt_int = gt_int.to(device)
+                
+                # Use scaled mesh renderer with appropriate near/far for smaller mesh
+                mesh_half_size = 0.5 / scale
+                renderer.rendering_options.near = max(0.01, distance - mesh_half_size - 0.5)
+                renderer.rendering_options.far = distance + mesh_half_size + 0.5
+                
+                try:
+                    gt_res = renderer.render(scaled_rep, gt_ext, gt_int)
+                    gt_view_images.append(gt_res['normal'])
+                except RuntimeError as e:
+                    print(f"[visualize_sample] Warning: GT view render failed for sample {i}: {e}")
+                    gt_view_images.append(torch.full((3, 512, 512), 0.5, device=device))
+        
+        result = {
+            'multiview': torch.stack(multiview_images),
+        }
+        
+        if has_gt_camera and len(gt_view_images) > 0:
+            result['gt_view'] = torch.stack(gt_view_images)
+            
+        return result
+    
+    
+class SLatShape(SLatShapeVisMixin, SLat):
+    """
+    structured latent for shape generation
+    
+    Args:
+        roots (str): path to the dataset
+        resolution (int): resolution of the shape
+        min_aesthetic_score (float): minimum aesthetic score
+        max_tokens (int): maximum number of tokens
+        latent_key (str): key of the latent to be used
+        normalization (dict): normalization stats
+        pretrained_slat_dec (str): name of the pretrained slat decoder
+        slat_dec_path (str): path to the slat decoder, if given, will override the pretrained_slat_dec
+        slat_dec_ckpt (str): name of the slat decoder checkpoint
+        skip_list (str, optional): path to a file containing sha256 hashes to skip
+        skip_aesthetic_score_datasets (list, optional): list of dataset names to skip aesthetic score check
+    """
+    def __init__(self,
+        roots: str,
+        *,
+        resolution: int,
+        min_aesthetic_score: float = 5.0,
+        max_tokens: int = 32768,
+        normalization: Optional[dict] = None,
+        pretrained_slat_dec: str = 'microsoft/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16',
+        slat_dec_path: Optional[str] = None,
+        slat_dec_ckpt: Optional[str] = None,
+        skip_list: Optional[str] = None,
+        skip_aesthetic_score_datasets: Optional[list] = None,
+    ):
+        super().__init__(
+            roots,
+            min_aesthetic_score=min_aesthetic_score,
+            max_tokens=max_tokens,
+            latent_key='shape_latent',
+            normalization=normalization,
+            pretrained_slat_dec=pretrained_slat_dec,
+            slat_dec_path=slat_dec_path,
+            slat_dec_ckpt=slat_dec_ckpt,
+            skip_list=skip_list,
+            skip_aesthetic_score_datasets=skip_aesthetic_score_datasets,
+        )
+        self.resolution = resolution
+
+
+class ImageConditionedSLatShape(ImageConditionedMixin, SLatShape):
+    """
+    Image conditioned structured latent for shape generation
+    """
+    pass
+
+
+class SLatShapeView(SLatShapeVisMixin, SLat):
+    """
+    View-based structured latent for shape generation.
+    
+    Data format: {sha256}/view{XX}.npz where each npz contains 'coords' and 'feats' keys.
+    
+    Args:
+        roots (str): path to the dataset
+        resolution (int): resolution of the shape
+        min_aesthetic_score (float): minimum aesthetic score
+        max_tokens (int): maximum number of tokens
+        num_views (int): Number of views to use (0 to num_views-1). Default is 2.
+        normalization (dict): normalization stats
+        pretrained_slat_dec (str): name of the pretrained slat decoder
+        slat_dec_path (str): path to the slat decoder, if given, will override the pretrained_slat_dec
+        slat_dec_ckpt (str): name of the slat decoder checkpoint
+        skip_list (str, optional): path to a file containing sha256 hashes to skip
+        skip_aesthetic_score_datasets (list, optional): list of dataset names to skip aesthetic score check
+    """
+    def __init__(self,
+        roots: str,
+        *,
+        resolution: int,
+        min_aesthetic_score: float = 5.0,
+        max_tokens: int = 32768,
+        num_views: int = 2,
+        normalization: Optional[dict] = None,
+        pretrained_slat_dec: str = 'microsoft/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16',
+        slat_dec_path: Optional[str] = None,
+        slat_dec_ckpt: Optional[str] = None,
+        skip_list: Optional[str] = None,
+        skip_aesthetic_score_datasets: Optional[list] = None,
+    ):
+        self.normalization = normalization
+        self.min_aesthetic_score = min_aesthetic_score
+        self.max_tokens = max_tokens
+        self.num_views = num_views
+        self.latent_key = 'shape_latent'
+        self.value_range = (0, 1)
+        
+        # Initialize parent with SLatVisMixin parameters
+        from .components import StandardDatasetBase
+        SLatVisMixin.__init__(
+            self,
+            roots,
+            pretrained_slat_dec=pretrained_slat_dec,
+            slat_dec_path=slat_dec_path,
+            slat_dec_ckpt=slat_dec_ckpt,
+        )
+        StandardDatasetBase.__init__(self, roots, skip_list=skip_list, skip_aesthetic_score_datasets=skip_aesthetic_score_datasets)
+        
+        self.resolution = resolution
+        
+        # Calculate loads for load balancing
+        self.loads = []
+        for _, sha256, _ in self.instances:
+            if f'{self.latent_key}_tokens' in self.metadata.columns:
+                try:
+                    self.loads.append(self.metadata.loc[sha256, f'{self.latent_key}_tokens'])
+                except:
+                    self.loads.append(self.max_tokens)
+            else:
+                self.loads.append(self.max_tokens)
+        
+        if self.normalization is not None:
+            self.mean = torch.tensor(self.normalization['mean']).reshape(1, -1)
+            self.std = torch.tensor(self.normalization['std']).reshape(1, -1)
+
+    def filter_metadata(self, metadata, dataset_name=None):
+        stats = {}
+        # View-based shape_latent uses columns like shape_latent_view00_encoded, shape_latent_view01_encoded, etc.
+        required_view_cols = [f'shape_latent_view{i:02d}_encoded' for i in range(self.num_views)]
+        existing_view_cols = [col for col in required_view_cols if col in metadata.columns]
+        
+        if existing_view_cols:
+            # Filter rows where all required views are encoded
+            # Note: NaN should be treated as False, so use == True for explicit comparison
+            has_all_views = (metadata[existing_view_cols] == True).all(axis=1)
+            metadata = metadata[has_all_views]
+            stats[f'With {self.num_views} view latents'] = len(metadata)
+        else:
+            # Fallback: check shape_latent_encoded column
+            if f'{self.latent_key}_encoded' in metadata.columns:
+                metadata = metadata[metadata[f'{self.latent_key}_encoded'] == True]
+                stats['With latent'] = len(metadata)
+        
+        # Skip aesthetic score check for specified datasets (e.g., texverse) or if column doesn't exist
+        skip_aesthetic = (
+            (dataset_name and dataset_name.lower() in [d.lower() for d in self.skip_aesthetic_score_datasets]) or
+            ('aesthetic_score' not in metadata.columns)
+        )
+        if skip_aesthetic:
+            stats[f'Aesthetic score check skipped'] = len(metadata)
+        else:
+            metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]
+            stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)
+        
+        # Filter by max_tokens if column exists
+        tokens_col = f'{self.latent_key}_tokens'
+        if tokens_col in metadata.columns:
+            metadata = metadata[metadata[tokens_col] <= self.max_tokens]
+            stats[f'Num tokens <= {self.max_tokens}'] = len(metadata)
+        
+        return metadata, stats
+
+    def get_instance(self, root, instance):
+        # View-based format: directory with view{XX}.npz files
+        latent_dir = os.path.join(root[self.latent_key], instance)
+        
+        # Randomly select a view from the configured range
+        view_idx = np.random.randint(0, self.num_views)
+        view_file = f'view{view_idx:02d}.npz'
+        
+        # Store view info for ViewImageConditionedMixin
+        self._current_view_idx = view_idx
+        self._current_latent_dir = latent_dir
+        
+        data = np.load(os.path.join(latent_dir, view_file))
+        coords = torch.tensor(data['coords']).int()
+        feats = torch.tensor(data['feats']).float()
+        if self.normalization is not None:
+            feats = (feats - self.mean) / self.std
+        return {
+            'coords': coords,
+            'feats': feats,
+            'view_idx': view_idx,
+        }
+
+    @staticmethod
+    def collate_fn(batch, split_size=None):
+        if split_size is None:
+            group_idx = [list(range(len(batch)))]
+        else:
+            group_idx = load_balanced_group_indices([b['coords'].shape[0] for b in batch], split_size)
+        packs = []
+        for group in group_idx:
+            sub_batch = [batch[i] for i in group]
+            pack = {}
+            coords = []
+            feats = []
+            layout = []
+            start = 0
+            for i, b in enumerate(sub_batch):
+                coords.append(torch.cat([torch.full((b['coords'].shape[0], 1), i, dtype=torch.int32), b['coords']], dim=-1))
+                feats.append(b['feats'])
+                layout.append(slice(start, start + b['coords'].shape[0]))
+                start += b['coords'].shape[0]
+            coords = torch.cat(coords)
+            feats = torch.cat(feats)
+            pack['x_0'] = SparseTensor(
+                coords=coords,
+                feats=feats,
+            )
+            pack['x_0']._shape = torch.Size([len(group), *sub_batch[0]['feats'].shape[1:]])
+            pack['x_0'].register_spatial_cache('layout', layout)
+            
+            # collate other data
+            keys = [k for k in sub_batch[0].keys() if k not in ['coords', 'feats']]
+            for k in keys:
+                if isinstance(sub_batch[0][k], torch.Tensor):
+                    pack[k] = torch.stack([b[k] for b in sub_batch])
+                elif isinstance(sub_batch[0][k], list):
+                    pack[k] = sum([b[k] for b in sub_batch], [])
+                else:
+                    pack[k] = [b[k] for b in sub_batch]
+                    
+            packs.append(pack)
+          
+        if split_size is None:
+            return packs[0]
+        return packs
+
+
+class ViewImageConditionedSLatShapeView(ViewImageConditionedMixin, SLatShapeView):
+    """
+    Image-conditioned view-based structured latent for shape generation.
+    
+    Loads shape_latent from {sha256}/view{XX}.npz format and pairs with 
+    corresponding view from render_cond.
+    
+    Uses ViewImageConditionedMixin which reads mesh_scale from view{XX}_scale.json.
+    """
+    pass
diff --git a/pixal3d/datasets/structured_latent_svpbr.py b/pixal3d/datasets/structured_latent_svpbr.py
new file mode 100644
index 0000000..5aeb157
--- /dev/null
+++ b/pixal3d/datasets/structured_latent_svpbr.py
@@ -0,0 +1,666 @@
+import os
+os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1'
+import json
+from typing import *
+import numpy as np
+import torch
+import cv2
+import utils3d
+from .. import models
+from .components import StandardDatasetBase, ImageConditionedMixin, ViewImageConditionedMixin
+from ..modules.sparse import SparseTensor, sparse_cat
+from ..representations import MeshWithVoxel
+from ..renderers import PbrMeshRenderer, EnvMap
+from ..utils.data_utils import load_balanced_group_indices
+from ..utils.render_utils import yaw_pitch_r_fov_to_extrinsics_intrinsics
+
+
+class SLatPbrVisMixin:
+    def __init__(
+        self,
+        *args,
+        pretrained_pbr_slat_dec: str = 'JeffreyXiang/TRELLIS.2-4B/ckpts/tex_dec_next_dc_f16c32_fp16',
+        pbr_slat_dec_path: Optional[str] = None,
+        pbr_slat_dec_ckpt: Optional[str] = None,
+        pretrained_shape_slat_dec: str = 'JeffreyXiang/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16',
+        shape_slat_dec_path: Optional[str] = None,
+        shape_slat_dec_ckpt: Optional[str] = None,
+        **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        self.pbr_slat_dec = None
+        self.pretrained_pbr_slat_dec = pretrained_pbr_slat_dec
+        self.pbr_slat_dec_path = pbr_slat_dec_path
+        self.pbr_slat_dec_ckpt = pbr_slat_dec_ckpt
+        self.shape_slat_dec = None
+        self.pretrained_shape_slat_dec = pretrained_shape_slat_dec
+        self.shape_slat_dec_path = shape_slat_dec_path
+        self.shape_slat_dec_ckpt = shape_slat_dec_ckpt
+        
+    def _loading_slat_dec(self):
+        if self.pbr_slat_dec is not None and self.shape_slat_dec is not None:
+            return
+        if self.pbr_slat_dec_path is not None:
+            cfg = json.load(open(os.path.join(self.pbr_slat_dec_path, 'config.json'), 'r'))
+            decoder = getattr(models, cfg['models']['decoder']['name'])(**cfg['models']['decoder']['args'])
+            ckpt_path = os.path.join(self.pbr_slat_dec_path, 'ckpts', f'decoder_{self.pbr_slat_dec_ckpt}.pt')
+            decoder.load_state_dict(torch.load(ckpt_path, map_location='cpu', weights_only=True))
+        else:
+            decoder = models.from_pretrained(self.pretrained_pbr_slat_dec)
+        self.pbr_slat_dec = decoder.cuda().eval()
+
+        if self.shape_slat_dec_path is not None:
+            cfg = json.load(open(os.path.join(self.shape_slat_dec_path, 'config.json'), 'r'))
+            decoder = getattr(models, cfg['models']['decoder']['name'])(**cfg['models']['decoder']['args'])
+            ckpt_path = os.path.join(self.shape_slat_dec_path, 'ckpts', f'decoder_{self.shape_slat_dec_ckpt}.pt')
+            decoder.load_state_dict(torch.load(ckpt_path, map_location='cpu', weights_only=True))
+        else:
+            decoder = models.from_pretrained(self.pretrained_shape_slat_dec)
+        decoder.set_resolution(self.resolution)
+        self.shape_slat_dec = decoder.cuda().eval()
+
+    def _delete_slat_dec(self):
+        del self.pbr_slat_dec
+        self.pbr_slat_dec = None
+        del self.shape_slat_dec
+        self.shape_slat_dec = None
+        
+    @torch.no_grad()
+    def decode_latent(self, z, shape_z, batch_size=4):
+        self._loading_slat_dec()
+        reps = []
+        if self.shape_slat_normalization is not None:
+            shape_z = shape_z * self.shape_slat_std.to(z.device) + self.shape_slat_mean.to(z.device)
+        if self.pbr_slat_normalization is not None:
+            z = z * self.pbr_slat_std.to(z.device) + self.pbr_slat_mean.to(z.device)
+        for i in range(0, z.shape[0], batch_size):
+            mesh, subs = self.shape_slat_dec(shape_z[i:i+batch_size], return_subs=True)
+            vox = self.pbr_slat_dec(z[i:i+batch_size], guide_subs=subs) * 0.5 + 0.5
+            reps.extend([
+                MeshWithVoxel(
+                    m.vertices, m.faces,
+                    origin = [-0.5, -0.5, -0.5],
+                    voxel_size = 1 / self.resolution,
+                    coords = v.coords[:, 1:],
+                    attrs = v.feats,
+                    voxel_shape = torch.Size([*v.shape, *v.spatial_shape]),
+                    layout = self.layout,
+                )
+                for m, v in zip(mesh, vox)
+            ])
+        self._delete_slat_dec()
+        return reps
+    
+    @torch.no_grad()
+    def visualize_sample(self, sample: dict):
+        shape_z = sample['concat_cond'].cuda()
+        z = sample['x_0'].cuda()
+        reps = self.decode_latent(z, shape_z)
+        
+        # Extract camera parameters for GT view rendering (if available)
+        camera_angle_x = sample.get('camera_angle_x')
+        camera_distance = sample.get('camera_distance')
+        mesh_scale = sample.get('mesh_scale')
+        has_gt_camera = (
+            camera_angle_x is not None and
+            camera_distance is not None and
+            mesh_scale is not None
+        )
+        
+        # build camera
+        yaw = [0, np.pi/2, np.pi, 3*np.pi/2]
+        yaw_offset = -16 / 180 * np.pi
+        yaw = [y + yaw_offset for y in yaw]
+        pitch = [20 / 180 * np.pi for _ in range(4)]
+        exts, ints = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaw, pitch, 2, 30)
+        
+        # render
+        renderer = PbrMeshRenderer()
+        renderer.rendering_options.resolution = 512
+        renderer.rendering_options.near = 1
+        renderer.rendering_options.far = 100
+        renderer.rendering_options.ssaa = 2
+        renderer.rendering_options.peel_layers = 8
+        envmap = EnvMap(torch.tensor(
+            cv2.cvtColor(cv2.imread('assets/hdri/forest.exr', cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB),
+            dtype=torch.float32, device='cuda'
+        ))
+        
+        images = {}
+        gt_view_images = {}
+        for i, representation in enumerate(reps):
+            # Validate mesh data before rasterization (same as shape training)
+            verts = representation.vertices
+            faces = representation.faces
+            if verts.shape[0] == 0 or faces.shape[0] == 0:
+                print(f"[visualize_sample] Warning: sample {i} has empty mesh, skipping")
+                continue
+            if faces.max() >= verts.shape[0]:
+                print(f"[visualize_sample] Warning: sample {i} has out-of-bound face indices "
+                      f"(max face idx={faces.max().item()}, num verts={verts.shape[0]}), skipping")
+                continue
+            if torch.isnan(verts).any() or torch.isinf(verts).any():
+                print(f"[visualize_sample] Warning: sample {i} has NaN/Inf vertices, skipping")
+                continue
+
+            image = {}
+            tile = [2, 2]
+            try:
+                for j, (ext, intr) in enumerate(zip(exts, ints)):
+                    res = renderer.render(representation, ext, intr, envmap=envmap)
+                    for k, v in res.items():
+                        if k not in images:
+                            images[k] = []
+                        if k not in image:
+                            image[k] = torch.zeros(3, 1024, 1024).cuda()  
+                        image[k][:, 512 * (j // tile[1]):512 * (j // tile[1] + 1), 512 * (j % tile[1]):512 * (j % tile[1] + 1)] = v
+                for k in images.keys():
+                    images[k].append(image[k])
+            except RuntimeError as e:
+                print(f"[visualize_sample] Warning: render failed for sample {i}: {e}")
+                try:
+                    torch.cuda.synchronize()
+                except Exception:
+                    pass
+                try:
+                    torch.cuda.empty_cache()
+                except Exception:
+                    pass
+                continue
+            
+            # Render GT camera view
+            # Must scale mesh vertices by / mesh_scale to match ProjGrid's projection space.
+            # ProjGrid maps [-1,1]^3 -> / scale / 2 -> [-0.5/s, 0.5/s]^3
+            # Mesh vertices in [-0.5, 0.5]^3 -> / scale -> [-0.5/s, 0.5/s]^3 (equivalent)
+            if has_gt_camera:
+                try:
+                    scale = mesh_scale[i].item()
+                    distance = camera_distance[i].item()
+                    fov = camera_angle_x[i].item()
+                    device = representation.vertices.device
+                    
+                    # Scale mesh and voxel to match ProjGrid's projection space
+                    scaled_rep = MeshWithVoxel(
+                        vertices=representation.vertices / scale,
+                        faces=representation.faces,
+                        origin=(representation.origin / scale).tolist(),
+                        voxel_size=representation.voxel_size / scale,
+                        coords=representation.coords,
+                        attrs=representation.attrs,
+                        voxel_shape=representation.voxel_shape,
+                        layout=representation.layout,
+                    )
+                    
+                    cam_pos = torch.tensor([0.0, 0.0, distance], device=device)
+                    look_at = torch.tensor([0.0, 0.0, 0.0], device=device)
+                    cam_up = torch.tensor([0.0, 1.0, 0.0], device=device)
+                    
+                    gt_ext = utils3d.torch.extrinsics_look_at(cam_pos, look_at, cam_up)
+                    gt_int = utils3d.torch.intrinsics_from_fov_xy(
+                        torch.tensor(fov, device=device),
+                        torch.tensor(fov, device=device)
+                    )
+                    gt_ext = gt_ext.to(device)
+                    gt_int = gt_int.to(device)
+                    
+                    # Update near/far for the smaller scaled mesh
+                    mesh_half_size = 0.5 / scale
+                    renderer.rendering_options.near = max(0.01, distance - mesh_half_size - 0.5)
+                    renderer.rendering_options.far = distance + mesh_half_size + 0.5
+                    
+                    gt_res = renderer.render(scaled_rep, gt_ext, gt_int, envmap=envmap)
+                    for k, v in gt_res.items():
+                        gt_key = f'gt_view_{k}'
+                        if gt_key not in gt_view_images:
+                            gt_view_images[gt_key] = []
+                        gt_view_images[gt_key].append(v)
+                except RuntimeError as e:
+                    print(f"[visualize_sample] Warning: GT view render failed for sample {i}: {e}")
+                    try:
+                        torch.cuda.synchronize()
+                    except Exception:
+                        pass
+                    try:
+                        torch.cuda.empty_cache()
+                    except Exception:
+                        pass
+        
+        for k in images.keys():
+            images[k] = torch.stack(images[k], dim=0)
+        
+        for k, v in gt_view_images.items():
+            images[k] = torch.stack(v)
+        
+        return images
+    
+    
+class SLatPbr(SLatPbrVisMixin, StandardDatasetBase):
+    """
+    structured latent for sparse voxel pbr dataset
+    
+    Args:
+        roots (str): path to the dataset
+        latent_key (str): key of the latent to be used
+        min_aesthetic_score (float): minimum aesthetic score
+        normalization (dict): normalization stats
+        resolution (int): resolution of decoded sparse voxel
+        attrs (list): attributes to be decoded
+        pretained_slat_dec (str): name of the pretrained slat decoder
+        slat_dec_path (str): path to the slat decoder, if given, will override the pretrained_slat_dec
+        slat_dec_ckpt (str): name of the slat decoder checkpoint
+    """
+    def __init__(self,
+        roots: str,
+        *,
+        resolution: int,
+        min_aesthetic_score: float = 5.0,
+        max_tokens: int = 32768,
+        full_pbr: bool = False,
+        pbr_slat_normalization: Optional[dict] = None,
+        shape_slat_normalization: Optional[dict] = None,
+        attrs: list[str] = ['base_color', 'metallic', 'roughness', 'emissive', 'alpha'],
+        pretrained_pbr_slat_dec: str = 'JeffreyXiang/TRELLIS.2-4B/ckpts/tex_dec_next_dc_f16c32_fp16',
+        pbr_slat_dec_path: Optional[str] = None,
+        pbr_slat_dec_ckpt: Optional[str] = None,
+        pretrained_shape_slat_dec: str = 'JeffreyXiang/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16',
+        shape_slat_dec_path: Optional[str] = None,
+        shape_slat_dec_ckpt: Optional[str] = None,
+        **kwargs
+    ):  
+        self.resolution = resolution
+        self.pbr_slat_normalization = pbr_slat_normalization
+        self.shape_slat_normalization = shape_slat_normalization
+        self.min_aesthetic_score = min_aesthetic_score
+        self.max_tokens = max_tokens
+        self.full_pbr = full_pbr
+        self.value_range = (0, 1)
+        
+        super().__init__(
+            roots,
+            pretrained_pbr_slat_dec=pretrained_pbr_slat_dec,
+            pbr_slat_dec_path=pbr_slat_dec_path,
+            pbr_slat_dec_ckpt=pbr_slat_dec_ckpt,
+            pretrained_shape_slat_dec=pretrained_shape_slat_dec,
+            shape_slat_dec_path=shape_slat_dec_path,
+            shape_slat_dec_ckpt=shape_slat_dec_ckpt,
+            **kwargs
+        )
+        
+        self.loads = [self.metadata.loc[sha256, 'pbr_latent_tokens'] for _, sha256, _ in self.instances]
+        
+        if self.pbr_slat_normalization is not None:
+            self.pbr_slat_mean = torch.tensor(self.pbr_slat_normalization['mean']).reshape(1, -1)
+            self.pbr_slat_std = torch.tensor(self.pbr_slat_normalization['std']).reshape(1, -1)
+        
+        if self.shape_slat_normalization is not None:
+            self.shape_slat_mean = torch.tensor(self.shape_slat_normalization['mean']).reshape(1, -1)
+            self.shape_slat_std = torch.tensor(self.shape_slat_normalization['std']).reshape(1, -1)
+        
+        self.attrs = attrs
+        self.channels = {
+            'base_color': 3,
+            'metallic': 1,
+            'roughness': 1,
+            'emissive': 3,
+            'alpha': 1,
+        }
+        self.layout = {}
+        start = 0
+        for attr in attrs:
+            self.layout[attr] = slice(start, start + self.channels[attr])
+            start += self.channels[attr]
+            
+    def filter_metadata(self, metadata, dataset_name=None):
+        stats = {}
+        metadata = metadata[metadata['pbr_latent_encoded'] == True]
+        stats['With PBR latent'] = len(metadata)
+        metadata = metadata[metadata['shape_latent_encoded'] == True]
+        stats['With shape latent'] = len(metadata)
+        metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]
+        stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)
+        metadata = metadata[metadata['pbr_latent_tokens'] <= self.max_tokens]
+        stats[f'Num tokens <= {self.max_tokens}'] = len(metadata)
+        if self.full_pbr:
+            metadata = metadata[metadata['num_basecolor_tex'] > 0]
+            metadata = metadata[metadata['num_metallic_tex'] > 0]
+            metadata = metadata[metadata['num_roughness_tex'] > 0]
+            stats['Full PBR'] = len(metadata)
+        return metadata, stats
+    
+    def get_instance(self, root, instance):
+        # PBR latent
+        data = np.load(os.path.join(root['pbr_latent'], f'{instance}.npz'))
+        coords = torch.tensor(data['coords']).int()
+        coords = torch.cat([torch.zeros_like(coords)[:, :1], coords], dim=1)
+        feats = torch.tensor(data['feats']).float()
+        if self.pbr_slat_normalization is not None:
+            feats = (feats - self.pbr_slat_mean) / self.pbr_slat_std
+        pbr_z = SparseTensor(feats, coords)
+        
+        # Shape latent
+        data = np.load(os.path.join(root['shape_latent'], f'{instance}.npz'))
+        coords = torch.tensor(data['coords']).int()
+        coords = torch.cat([torch.zeros_like(coords)[:, :1], coords], dim=1)
+        feats = torch.tensor(data['feats']).float()
+        if self.shape_slat_normalization is not None:
+            feats = (feats - self.shape_slat_mean) / self.shape_slat_std
+        shape_z = SparseTensor(feats, coords)
+        
+        assert torch.equal(shape_z.coords, pbr_z.coords), \
+            f"Shape latent and PBR latent have different coordinates: {shape_z.coords.shape} vs {pbr_z.coords.shape}"
+            
+        return {
+            'x_0': pbr_z,
+            'concat_cond': shape_z,
+        }
+        
+    @staticmethod
+    def collate_fn(batch, split_size=None):
+        if split_size is None:
+            group_idx = [list(range(len(batch)))]
+        else:
+            group_idx = load_balanced_group_indices([b['x_0'].feats.shape[0] for b in batch], split_size)
+        packs = []
+        for group in group_idx:
+            sub_batch = [batch[i] for i in group]
+            pack = {}
+
+            keys = [k for k in sub_batch[0].keys()]
+            for k in keys:
+                if isinstance(sub_batch[0][k], torch.Tensor):
+                    pack[k] = torch.stack([b[k] for b in sub_batch])
+                elif isinstance(sub_batch[0][k], SparseTensor):
+                    pack[k] = sparse_cat([b[k] for b in sub_batch], dim=0)
+                elif isinstance(sub_batch[0][k], list):
+                    pack[k] = sum([b[k] for b in sub_batch], [])
+                else:
+                    pack[k] = [b[k] for b in sub_batch]
+            
+            packs.append(pack)
+        
+        if split_size is None:
+            return packs[0]
+        return packs
+
+
+class ImageConditionedSLatPbr(ImageConditionedMixin, SLatPbr):
+    """
+    Image conditioned structured latent dataset
+    """
+    pass
+
+
+class SLatPbrView(SLatPbrVisMixin, StandardDatasetBase):
+    """
+    View-based structured latent for PBR/texture generation with view-aligned projection.
+    
+    Data format: 
+        PBR latent:   {sha256}/view{XX}.npz  (coords + feats)
+        Shape latent: {sha256}/view{XX}.npz  (coords + feats, from shape_latent_view dir)
+    
+    Each view's PBR latent and Shape latent share the same sparse coordinates.
+    
+    Args:
+        roots (str): path to the dataset
+        resolution (int): resolution of decoded sparse voxel
+        min_aesthetic_score (float): minimum aesthetic score
+        max_tokens (int): maximum number of tokens
+        num_views (int): Number of views to use (0 to num_views-1). Default is 2.
+        full_pbr (bool): Whether to require full PBR textures
+        pbr_slat_normalization (dict): normalization stats for PBR latent
+        shape_slat_normalization (dict): normalization stats for shape latent
+        attrs (list): PBR attributes to decode
+        pretrained_pbr_slat_dec (str): pretrained PBR decoder name
+        pretrained_shape_slat_dec (str): pretrained shape decoder name
+        skip_list (str, optional): path to a file containing sha256 hashes to skip
+        skip_aesthetic_score_datasets (list, optional): datasets to skip aesthetic score check
+    """
+    def __init__(self,
+        roots: str,
+        *,
+        resolution: int,
+        min_aesthetic_score: float = 5.0,
+        max_tokens: int = 32768,
+        num_views: int = 2,
+        full_pbr: bool = False,
+        pbr_slat_normalization: Optional[dict] = None,
+        shape_slat_normalization: Optional[dict] = None,
+        attrs: list[str] = ['base_color', 'metallic', 'roughness', 'emissive', 'alpha'],
+        pretrained_pbr_slat_dec: str = 'microsoft/TRELLIS.2-4B/ckpts/tex_dec_next_dc_f16c32_fp16',
+        pbr_slat_dec_path: Optional[str] = None,
+        pbr_slat_dec_ckpt: Optional[str] = None,
+        pretrained_shape_slat_dec: str = 'microsoft/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16',
+        shape_slat_dec_path: Optional[str] = None,
+        shape_slat_dec_ckpt: Optional[str] = None,
+        skip_list: Optional[str] = None,
+        skip_aesthetic_score_datasets: Optional[list] = None,
+    ):
+        self.resolution = resolution
+        self.pbr_slat_normalization = pbr_slat_normalization
+        self.shape_slat_normalization = shape_slat_normalization
+        self.min_aesthetic_score = min_aesthetic_score
+        self.max_tokens = max_tokens
+        self.num_views = num_views
+        self.full_pbr = full_pbr
+        self.value_range = (0, 1)
+        self.skip_aesthetic_score_datasets = set(skip_aesthetic_score_datasets or [])
+        
+        # Initialize visualization mixin
+        SLatPbrVisMixin.__init__(
+            self,
+            roots,
+            pretrained_pbr_slat_dec=pretrained_pbr_slat_dec,
+            pbr_slat_dec_path=pbr_slat_dec_path,
+            pbr_slat_dec_ckpt=pbr_slat_dec_ckpt,
+            pretrained_shape_slat_dec=pretrained_shape_slat_dec,
+            shape_slat_dec_path=shape_slat_dec_path,
+            shape_slat_dec_ckpt=shape_slat_dec_ckpt,
+        )
+        StandardDatasetBase.__init__(
+            self, roots,
+            skip_list=skip_list,
+            skip_aesthetic_score_datasets=skip_aesthetic_score_datasets,
+        )
+        
+        # Calculate loads for load balancing
+        self.loads = []
+        for _, sha256, _ in self.instances:
+            if 'pbr_latent_tokens' in self.metadata.columns:
+                try:
+                    self.loads.append(self.metadata.loc[sha256, 'pbr_latent_tokens'])
+                except:
+                    self.loads.append(self.max_tokens)
+            else:
+                self.loads.append(self.max_tokens)
+        
+        if self.pbr_slat_normalization is not None:
+            self.pbr_slat_mean = torch.tensor(self.pbr_slat_normalization['mean']).reshape(1, -1)
+            self.pbr_slat_std = torch.tensor(self.pbr_slat_normalization['std']).reshape(1, -1)
+        
+        if self.shape_slat_normalization is not None:
+            self.shape_slat_mean = torch.tensor(self.shape_slat_normalization['mean']).reshape(1, -1)
+            self.shape_slat_std = torch.tensor(self.shape_slat_normalization['std']).reshape(1, -1)
+        
+        self.attrs = attrs
+        self.channels = {
+            'base_color': 3,
+            'metallic': 1,
+            'roughness': 1,
+            'emissive': 3,
+            'alpha': 1,
+        }
+        self.layout = {}
+        start = 0
+        for attr in attrs:
+            self.layout[attr] = slice(start, start + self.channels[attr])
+            start += self.channels[attr]
+
+    def filter_metadata(self, metadata, dataset_name=None):
+        stats = {}
+        # View-based PBR latent uses columns like pbr_latent_view00_encoded, etc.
+        required_pbr_view_cols = [f'pbr_latent_view{i:02d}_encoded' for i in range(self.num_views)]
+        existing_pbr_view_cols = [col for col in required_pbr_view_cols if col in metadata.columns]
+        
+        if existing_pbr_view_cols:
+            has_all_pbr_views = (metadata[existing_pbr_view_cols] == True).all(axis=1)
+            metadata = metadata[has_all_pbr_views]
+            stats[f'With {self.num_views} PBR view latents'] = len(metadata)
+        else:
+            # Fallback: check pbr_latent_encoded
+            if 'pbr_latent_encoded' in metadata.columns:
+                metadata = metadata[metadata['pbr_latent_encoded'] == True]
+                stats['With PBR latent'] = len(metadata)
+        
+        # Also require shape latent views
+        required_shape_view_cols = [f'shape_latent_view{i:02d}_encoded' for i in range(self.num_views)]
+        existing_shape_view_cols = [col for col in required_shape_view_cols if col in metadata.columns]
+        
+        if existing_shape_view_cols:
+            has_all_shape_views = (metadata[existing_shape_view_cols] == True).all(axis=1)
+            metadata = metadata[has_all_shape_views]
+            stats[f'With {self.num_views} shape view latents'] = len(metadata)
+        else:
+            if 'shape_latent_encoded' in metadata.columns:
+                metadata = metadata[metadata['shape_latent_encoded'] == True]
+                stats['With shape latent'] = len(metadata)
+        
+        # Skip aesthetic score check for specified datasets
+        skip_aesthetic = (
+            (dataset_name and dataset_name.lower() in [d.lower() for d in self.skip_aesthetic_score_datasets]) or
+            ('aesthetic_score' not in metadata.columns)
+        )
+        if skip_aesthetic:
+            stats[f'Aesthetic score check skipped'] = len(metadata)
+        else:
+            metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]
+            stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)
+        
+        # Filter by max_tokens if column exists
+        if 'pbr_latent_tokens' in metadata.columns:
+            metadata = metadata[metadata['pbr_latent_tokens'] <= self.max_tokens]
+            stats[f'Num tokens <= {self.max_tokens}'] = len(metadata)
+        
+        if self.full_pbr:
+            if 'num_basecolor_tex' in metadata.columns:
+                metadata = metadata[metadata['num_basecolor_tex'] > 0]
+            if 'num_metallic_tex' in metadata.columns:
+                metadata = metadata[metadata['num_metallic_tex'] > 0]
+            if 'num_roughness_tex' in metadata.columns:
+                metadata = metadata[metadata['num_roughness_tex'] > 0]
+            stats['Full PBR'] = len(metadata)
+        
+        return metadata, stats
+
+    def get_instance(self, root, instance):
+        # Randomly select a view from the configured range
+        view_idx = np.random.randint(0, self.num_views)
+        view_file = f'view{view_idx:02d}.npz'
+        
+        # Store view info for ViewImageConditionedMixin
+        self._current_view_idx = view_idx
+        
+        # Load PBR latent for this view
+        pbr_latent_dir = os.path.join(root['pbr_latent'], instance)
+        self._current_latent_dir = pbr_latent_dir
+        
+        data = np.load(os.path.join(pbr_latent_dir, view_file))
+        pbr_coords = torch.tensor(data['coords']).int()
+        pbr_feats = torch.tensor(data['feats']).float()
+        if self.pbr_slat_normalization is not None:
+            pbr_feats = (pbr_feats - self.pbr_slat_mean) / self.pbr_slat_std
+        
+        # Load Shape latent for this view (as concat_cond)
+        shape_latent_dir = os.path.join(root['shape_latent'], instance)
+        data = np.load(os.path.join(shape_latent_dir, view_file))
+        shape_coords = torch.tensor(data['coords']).int()
+        shape_feats = torch.tensor(data['feats']).float()
+        if self.shape_slat_normalization is not None:
+            shape_feats = (shape_feats - self.shape_slat_mean) / self.shape_slat_std
+        
+        # Verify coordinates match
+        assert torch.equal(pbr_coords, shape_coords), \
+            f"PBR and shape latent coordinates mismatch for {instance}/view{view_idx:02d}"
+        
+        return {
+            'coords': pbr_coords,
+            'pbr_feats': pbr_feats,
+            'shape_feats': shape_feats,
+            'view_idx': view_idx,
+        }
+
+    @staticmethod
+    def collate_fn(batch, split_size=None):
+        if split_size is None:
+            group_idx = [list(range(len(batch)))]
+        else:
+            group_idx = load_balanced_group_indices([b['coords'].shape[0] for b in batch], split_size)
+        packs = []
+        for group in group_idx:
+            sub_batch = [batch[i] for i in group]
+            pack = {}
+            
+            # Build x_0 (PBR latent) and concat_cond (shape latent) as SparseTensors
+            coords_list = []
+            pbr_feats_list = []
+            shape_feats_list = []
+            layout = []
+            start = 0
+            for i, b in enumerate(sub_batch):
+                batch_coords = torch.cat([
+                    torch.full((b['coords'].shape[0], 1), i, dtype=torch.int32),
+                    b['coords']
+                ], dim=-1)
+                coords_list.append(batch_coords)
+                pbr_feats_list.append(b['pbr_feats'])
+                shape_feats_list.append(b['shape_feats'])
+                layout.append(slice(start, start + b['coords'].shape[0]))
+                start += b['coords'].shape[0]
+            
+            all_coords = torch.cat(coords_list)
+            
+            # x_0: PBR latent
+            pack['x_0'] = SparseTensor(
+                coords=all_coords,
+                feats=torch.cat(pbr_feats_list),
+            )
+            pack['x_0']._shape = torch.Size([len(group), *sub_batch[0]['pbr_feats'].shape[1:]])
+            pack['x_0'].register_spatial_cache('layout', layout)
+            
+            # concat_cond: Shape latent (same coordinates)
+            pack['concat_cond'] = SparseTensor(
+                coords=all_coords.clone(),
+                feats=torch.cat(shape_feats_list),
+            )
+            pack['concat_cond']._shape = torch.Size([len(group), *sub_batch[0]['shape_feats'].shape[1:]])
+            pack['concat_cond'].register_spatial_cache('layout', layout)
+            
+            # collate other data (excluding already handled fields)
+            skip_keys = {'coords', 'pbr_feats', 'shape_feats'}
+            keys = [k for k in sub_batch[0].keys() if k not in skip_keys]
+            for k in keys:
+                if isinstance(sub_batch[0][k], torch.Tensor):
+                    pack[k] = torch.stack([b[k] for b in sub_batch])
+                elif isinstance(sub_batch[0][k], list):
+                    pack[k] = sum([b[k] for b in sub_batch], [])
+                else:
+                    pack[k] = [b[k] for b in sub_batch]
+            
+            packs.append(pack)
+        
+        if split_size is None:
+            return packs[0]
+        return packs
+
+
+class ViewImageConditionedSLatPbrView(ViewImageConditionedMixin, SLatPbrView):
+    """
+    Image-conditioned view-based structured latent for PBR/texture generation
+    with view-aligned projection.
+    
+    Loads PBR latent and shape latent from {sha256}/view{XX}.npz format and pairs
+    with corresponding view from render_cond.
+    
+    Uses ViewImageConditionedMixin which reads mesh_scale from view{XX}_scale.json
+    and provides camera parameters for 3D-to-2D projection.
+    """
+    pass
diff --git a/pixal3d/models/__init__.py b/pixal3d/models/__init__.py
new file mode 100644
index 0000000..d4fed03
--- /dev/null
+++ b/pixal3d/models/__init__.py
@@ -0,0 +1,78 @@
+import importlib
+
+__attributes = {
+    # Sparse Structure
+    'SparseStructureEncoder': 'sparse_structure_vae',
+    'SparseStructureDecoder': 'sparse_structure_vae',
+    'SparseStructureFlowModel': 'sparse_structure_flow',
+    
+    # SLat Generation
+    'SLatFlowModel': 'structured_latent_flow',
+    'ElasticSLatFlowModel': 'structured_latent_flow',
+    
+    # SC-VAEs
+    'SparseUnetVaeEncoder': 'sc_vaes.sparse_unet_vae',
+    'SparseUnetVaeDecoder': 'sc_vaes.sparse_unet_vae',
+    'FlexiDualGridVaeEncoder': 'sc_vaes.fdg_vae',
+    'FlexiDualGridVaeDecoder': 'sc_vaes.fdg_vae'
+}
+
+__submodules = []
+
+__all__ = list(__attributes.keys()) + __submodules
+
+def __getattr__(name):
+    if name not in globals():
+        if name in __attributes:
+            module_name = __attributes[name]
+            module = importlib.import_module(f".{module_name}", __name__)
+            globals()[name] = getattr(module, name)
+        elif name in __submodules:
+            module = importlib.import_module(f".{name}", __name__)
+            globals()[name] = module
+        else:
+            raise AttributeError(f"module {__name__} has no attribute {name}")
+    return globals()[name]
+
+
+def from_pretrained(path: str, **kwargs):
+    """
+    Load a model from a pretrained checkpoint.
+
+    Args:
+        path: The path to the checkpoint. Can be either local path or a Hugging Face model name.
+              NOTE: config file and model file should take the name f'{path}.json' and f'{path}.safetensors' respectively.
+        **kwargs: Additional arguments for the model constructor.
+    """
+    import os
+    import json
+    from safetensors.torch import load_file
+    is_local = os.path.exists(f"{path}.json") and os.path.exists(f"{path}.safetensors")
+
+    if is_local:
+        config_file = f"{path}.json"
+        model_file = f"{path}.safetensors"
+    else:
+        from huggingface_hub import hf_hub_download
+        path_parts = path.split('/')
+        repo_id = f'{path_parts[0]}/{path_parts[1]}'
+        model_name = '/'.join(path_parts[2:])
+        config_file = hf_hub_download(repo_id, f"{model_name}.json")
+        model_file = hf_hub_download(repo_id, f"{model_name}.safetensors")
+
+    with open(config_file, 'r') as f:
+        config = json.load(f)
+    model = __getattr__(config['name'])(**config['args'], **kwargs)
+    model.load_state_dict(load_file(model_file), strict=False)
+
+    return model
+
+
+# For Pylance
+if __name__ == '__main__':
+    from .sparse_structure_vae import SparseStructureEncoder, SparseStructureDecoder
+    from .sparse_structure_flow import SparseStructureFlowModel
+    from .structured_latent_flow import SLatFlowModel, ElasticSLatFlowModel
+        
+    from .sc_vaes.sparse_unet_vae import SparseUnetVaeEncoder, SparseUnetVaeDecoder
+    from .sc_vaes.fdg_vae import FlexiDualGridVaeEncoder, FlexiDualGridVaeDecoder
diff --git a/pixal3d/models/sc_vaes/fdg_vae.py b/pixal3d/models/sc_vaes/fdg_vae.py
new file mode 100644
index 0000000..0209e69
--- /dev/null
+++ b/pixal3d/models/sc_vaes/fdg_vae.py
@@ -0,0 +1,110 @@
+from typing import *
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from ...modules import sparse as sp
+from .sparse_unet_vae import (
+    SparseResBlock3d,
+    SparseConvNeXtBlock3d,
+    
+    SparseResBlockDownsample3d,
+    SparseResBlockUpsample3d,
+    SparseResBlockS2C3d,
+    SparseResBlockC2S3d,
+)
+from .sparse_unet_vae import (
+    SparseUnetVaeEncoder,
+    SparseUnetVaeDecoder,
+)
+from ...representations import Mesh
+from o_voxel.convert import flexible_dual_grid_to_mesh
+
+
+class FlexiDualGridVaeEncoder(SparseUnetVaeEncoder):
+    def __init__(
+        self,
+        model_channels: List[int],
+        latent_channels: int,
+        num_blocks: List[int],
+        block_type: List[str],
+        down_block_type: List[str],
+        block_args: List[Dict[str, Any]],
+        use_fp16: bool = False,
+    ):
+        super().__init__(
+            6,
+            model_channels,
+            latent_channels,
+            num_blocks,
+            block_type,
+            down_block_type,
+            block_args,
+            use_fp16,
+        )
+        
+    def forward(self, vertices: sp.SparseTensor, intersected: sp.SparseTensor, sample_posterior=False, return_raw=False):
+        x = vertices.replace(torch.cat([
+            vertices.feats - 0.5,
+            intersected.feats.float() - 0.5,
+        ], dim=1))
+        return super().forward(x, sample_posterior, return_raw)
+    
+    
+class FlexiDualGridVaeDecoder(SparseUnetVaeDecoder):
+    def __init__(
+        self,
+        resolution: int,
+        model_channels: List[int],
+        latent_channels: int,
+        num_blocks: List[int],
+        block_type: List[str],
+        up_block_type: List[str],
+        block_args: List[Dict[str, Any]],
+        voxel_margin: float = 0.5,
+        use_fp16: bool = False,
+    ):
+        self.resolution = resolution
+        self.voxel_margin = voxel_margin
+        
+        super().__init__(
+            7,
+            model_channels,
+            latent_channels,
+            num_blocks,
+            block_type,
+            up_block_type,
+            block_args,
+            use_fp16,
+        )
+
+    def set_resolution(self, resolution: int) -> None:
+        self.resolution = resolution
+        
+    def forward(self, x: sp.SparseTensor, gt_intersected: sp.SparseTensor = None, **kwargs):
+        decoded = super().forward(x, **kwargs)
+        if self.training:
+            h, subs_gt, subs = decoded
+            vertices = h.replace((1 + 2 * self.voxel_margin) * F.sigmoid(h.feats[..., 0:3]) - self.voxel_margin)
+            intersected_logits = h.replace(h.feats[..., 3:6])
+            quad_lerp = h.replace(F.softplus(h.feats[..., 6:7]))
+            mesh = [Mesh(*flexible_dual_grid_to_mesh(
+                v.coords[:, 1:], v.feats, i.feats, q.feats,
+                aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
+                grid_size=self.resolution,
+                train=True
+            )) for v, i, q in zip(vertices, gt_intersected, quad_lerp)]
+            return mesh, vertices, intersected_logits, subs_gt, subs
+        else:
+            out_list = list(decoded) if isinstance(decoded, tuple) else [decoded]
+            h = out_list[0]
+            vertices = h.replace((1 + 2 * self.voxel_margin) * F.sigmoid(h.feats[..., 0:3]) - self.voxel_margin)
+            intersected = h.replace(h.feats[..., 3:6] > 0)
+            quad_lerp = h.replace(F.softplus(h.feats[..., 6:7]))
+            mesh = [Mesh(*flexible_dual_grid_to_mesh(
+                v.coords[:, 1:], v.feats, i.feats, q.feats,
+                aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
+                grid_size=self.resolution,
+                train=False
+            )) for v, i, q in zip(vertices, intersected, quad_lerp)]
+            out_list[0] = mesh
+            return out_list[0] if len(out_list) == 1 else tuple(out_list)
diff --git a/pixal3d/models/sc_vaes/sparse_unet_vae.py b/pixal3d/models/sc_vaes/sparse_unet_vae.py
new file mode 100644
index 0000000..b9902a1
--- /dev/null
+++ b/pixal3d/models/sc_vaes/sparse_unet_vae.py
@@ -0,0 +1,522 @@
+from typing import *
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from ...modules.utils import convert_module_to_f16, convert_module_to_f32, zero_module
+from ...modules import sparse as sp
+from ...modules.norm import LayerNorm32
+
+
+class SparseResBlock3d(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        out_channels: Optional[int] = None,
+        downsample: bool = False,
+        upsample: bool = False,
+        resample_mode: Literal['nearest', 'spatial2channel'] = 'nearest',
+        use_checkpoint: bool = False,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.downsample = downsample
+        self.upsample = upsample
+        self.resample_mode = resample_mode
+        self.use_checkpoint = use_checkpoint
+        
+        assert not (downsample and upsample), "Cannot downsample and upsample at the same time"
+
+        self.norm1 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
+        self.norm2 = LayerNorm32(self.out_channels, elementwise_affine=False, eps=1e-6)
+        if resample_mode == 'nearest':
+            self.conv1 = sp.SparseConv3d(channels, self.out_channels, 3)
+        elif resample_mode =='spatial2channel' and not self.downsample:
+            self.conv1 = sp.SparseConv3d(channels, self.out_channels * 8, 3)
+        elif resample_mode =='spatial2channel' and self.downsample:
+            self.conv1 = sp.SparseConv3d(channels, self.out_channels // 8, 3)
+        self.conv2 = zero_module(sp.SparseConv3d(self.out_channels, self.out_channels, 3))
+        if resample_mode == 'nearest':
+            self.skip_connection = sp.SparseLinear(channels, self.out_channels) if channels != self.out_channels else nn.Identity()
+        elif resample_mode =='spatial2channel' and self.downsample:
+            self.skip_connection = lambda x: x.replace(x.feats.reshape(x.feats.shape[0], out_channels, channels * 8 // out_channels).mean(dim=-1))
+        elif resample_mode =='spatial2channel' and not self.downsample:
+            self.skip_connection = lambda x: x.replace(x.feats.repeat_interleave(out_channels // (channels // 8), dim=1))
+        self.updown = None
+        if self.downsample:
+            if resample_mode == 'nearest':
+                self.updown = sp.SparseDownsample(2)
+            elif resample_mode =='spatial2channel':
+                self.updown = sp.SparseSpatial2Channel(2)
+        elif self.upsample:
+            self.to_subdiv = sp.SparseLinear(channels, 8)
+            if resample_mode == 'nearest':
+                self.updown = sp.SparseUpsample(2)
+            elif resample_mode =='spatial2channel':
+                self.updown = sp.SparseChannel2Spatial(2)
+
+    def _updown(self, x: sp.SparseTensor, subdiv: sp.SparseTensor = None) -> sp.SparseTensor:
+        if self.downsample:
+            x = self.updown(x)
+        elif self.upsample:
+            x = self.updown(x, subdiv.replace(subdiv.feats > 0))
+        return x
+
+    def _forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
+        subdiv = None
+        if self.upsample:
+            subdiv = self.to_subdiv(x)
+        h = x.replace(self.norm1(x.feats))
+        h = h.replace(F.silu(h.feats))
+        if self.resample_mode == 'spatial2channel':
+            h = self.conv1(h)
+        h = self._updown(h, subdiv)
+        x = self._updown(x, subdiv)
+        if self.resample_mode == 'nearest':
+            h = self.conv1(h)
+        h = h.replace(self.norm2(h.feats))
+        h = h.replace(F.silu(h.feats))
+        h = self.conv2(h)
+        h = h + self.skip_connection(x)
+        if self.upsample:
+            return h, subdiv
+        return h
+    
+    def forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
+        if self.use_checkpoint:
+            return torch.utils.checkpoint.checkpoint(self._forward, x, use_reentrant=False)
+        else:
+            return self._forward(x)
+
+
+class SparseResBlockDownsample3d(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        out_channels: Optional[int] = None,
+        use_checkpoint: bool = False,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_checkpoint = use_checkpoint
+        
+        self.norm1 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
+        self.norm2 = LayerNorm32(self.out_channels, elementwise_affine=False, eps=1e-6)
+        self.conv1 = sp.SparseConv3d(channels, self.out_channels, 3)
+        self.conv2 = zero_module(sp.SparseConv3d(self.out_channels, self.out_channels, 3))
+        self.skip_connection = sp.SparseLinear(channels, self.out_channels) if channels != self.out_channels else nn.Identity()
+        self.updown = sp.SparseDownsample(2)
+
+    def _forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
+        h = x.replace(self.norm1(x.feats))
+        h = h.replace(F.silu(h.feats))
+        h = self.updown(h)
+        x = self.updown(x)
+        h = self.conv1(h)
+        h = h.replace(self.norm2(h.feats))
+        h = h.replace(F.silu(h.feats))
+        h = self.conv2(h)
+        h = h + self.skip_connection(x)
+        return h
+    
+    def forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
+        if self.use_checkpoint:
+            return torch.utils.checkpoint.checkpoint(self._forward, x, use_reentrant=False)
+        else:
+            return self._forward(x)
+
+
+class SparseResBlockUpsample3d(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        out_channels: Optional[int] = None,
+        use_checkpoint: bool = False,
+        pred_subdiv: bool = True,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_checkpoint = use_checkpoint
+        self.pred_subdiv = pred_subdiv
+        
+        self.norm1 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
+        self.norm2 = LayerNorm32(self.out_channels, elementwise_affine=False, eps=1e-6)
+        self.conv1 = sp.SparseConv3d(channels, self.out_channels, 3)
+        self.conv2 = zero_module(sp.SparseConv3d(self.out_channels, self.out_channels, 3))
+        self.skip_connection = sp.SparseLinear(channels, self.out_channels) if channels != self.out_channels else nn.Identity()
+        if self.pred_subdiv:
+            self.to_subdiv = sp.SparseLinear(channels, 8)
+        self.updown = sp.SparseUpsample(2)
+
+    def _forward(self, x: sp.SparseTensor, subdiv: sp.SparseTensor = None) -> sp.SparseTensor:
+        if self.pred_subdiv:
+            subdiv = self.to_subdiv(x)
+        h = x.replace(self.norm1(x.feats))
+        h = h.replace(F.silu(h.feats))
+        subdiv_binarized = subdiv.replace(subdiv.feats > 0) if subdiv is not None else None
+        h = self.updown(h, subdiv_binarized)
+        x = self.updown(x, subdiv_binarized)
+        h = self.conv1(h)
+        h = h.replace(self.norm2(h.feats))
+        h = h.replace(F.silu(h.feats))
+        h = self.conv2(h)
+        h = h + self.skip_connection(x)
+        if self.pred_subdiv:
+            return h, subdiv
+        else:
+            return h
+    
+    def forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
+        if self.use_checkpoint:
+            return torch.utils.checkpoint.checkpoint(self._forward, x, use_reentrant=False)
+        else:
+            return self._forward(x)
+
+
+class SparseResBlockS2C3d(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        out_channels: Optional[int] = None,
+        use_checkpoint: bool = False,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_checkpoint = use_checkpoint
+        
+        self.norm1 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
+        self.norm2 = LayerNorm32(self.out_channels, elementwise_affine=False, eps=1e-6)
+        self.conv1 = sp.SparseConv3d(channels, self.out_channels // 8, 3)
+        self.conv2 = zero_module(sp.SparseConv3d(self.out_channels, self.out_channels, 3))
+        self.skip_connection = lambda x: x.replace(x.feats.reshape(x.feats.shape[0], out_channels, channels * 8 // out_channels).mean(dim=-1))
+        self.updown = sp.SparseSpatial2Channel(2)
+
+    def _forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
+        h = x.replace(self.norm1(x.feats))
+        h = h.replace(F.silu(h.feats))
+        h = self.conv1(h)
+        h = self.updown(h)
+        x = self.updown(x)
+        h = h.replace(self.norm2(h.feats))
+        h = h.replace(F.silu(h.feats))
+        h = self.conv2(h)
+        h = h + self.skip_connection(x)
+        return h
+    
+    def forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
+        if self.use_checkpoint:
+            return torch.utils.checkpoint.checkpoint(self._forward, x, use_reentrant=False)
+        else:
+            return self._forward(x)
+
+
+class SparseResBlockC2S3d(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        out_channels: Optional[int] = None,
+        use_checkpoint: bool = False,
+        pred_subdiv: bool = True,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_checkpoint = use_checkpoint
+        self.pred_subdiv = pred_subdiv
+        
+        self.norm1 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
+        self.norm2 = LayerNorm32(self.out_channels, elementwise_affine=False, eps=1e-6)
+        self.conv1 = sp.SparseConv3d(channels, self.out_channels * 8, 3)
+        self.conv2 = zero_module(sp.SparseConv3d(self.out_channels, self.out_channels, 3))
+        self.skip_connection = lambda x: x.replace(x.feats.repeat_interleave(out_channels // (channels // 8), dim=1))
+        if pred_subdiv:
+            self.to_subdiv = sp.SparseLinear(channels, 8)
+        self.updown = sp.SparseChannel2Spatial(2)
+
+    def _forward(self, x: sp.SparseTensor, subdiv: sp.SparseTensor = None) -> sp.SparseTensor:
+        if self.pred_subdiv:
+            subdiv = self.to_subdiv(x)
+        h = x.replace(self.norm1(x.feats))
+        h = h.replace(F.silu(h.feats))
+        h = self.conv1(h)
+        subdiv_binarized = subdiv.replace(subdiv.feats > 0) if subdiv is not None else None
+        h = self.updown(h, subdiv_binarized)
+        x = self.updown(x, subdiv_binarized)
+        h = h.replace(self.norm2(h.feats))
+        h = h.replace(F.silu(h.feats))
+        h = self.conv2(h)
+        h = h + self.skip_connection(x)
+        if self.pred_subdiv:
+            return h, subdiv
+        else:
+            return h
+    
+    def forward(self, x: sp.SparseTensor, subdiv: sp.SparseTensor = None) -> sp.SparseTensor:
+        if self.use_checkpoint:
+            return torch.utils.checkpoint.checkpoint(self._forward, x, subdiv, use_reentrant=False)
+        else:
+            return self._forward(x, subdiv)
+        
+    
+class SparseConvNeXtBlock3d(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        mlp_ratio: float = 4.0,
+        use_checkpoint: bool = False,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.use_checkpoint = use_checkpoint
+        
+        self.norm = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
+        self.conv = sp.SparseConv3d(channels, channels, 3)
+        self.mlp = nn.Sequential(
+            nn.Linear(channels, int(channels * mlp_ratio)),
+            nn.SiLU(),
+            zero_module(nn.Linear(int(channels * mlp_ratio), channels)),
+        )
+
+    def _forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
+        h = self.conv(x)
+        h = h.replace(self.norm(h.feats))
+        h = h.replace(self.mlp(h.feats))
+        return h + x
+    
+    def forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
+        if self.use_checkpoint:
+            return torch.utils.checkpoint.checkpoint(self._forward, x, use_reentrant=False)
+        else:
+            return self._forward(x)
+
+
+class SparseUnetVaeEncoder(nn.Module):
+    """
+    Sparse Swin Transformer Unet VAE model.
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        model_channels: List[int],
+        latent_channels: int,
+        num_blocks: List[int],
+        block_type: List[str],
+        down_block_type: List[str],
+        block_args: List[Dict[str, Any]],
+        use_fp16: bool = False,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.num_blocks = num_blocks
+        self.dtype = torch.float16 if use_fp16 else torch.float32
+        self.dtype = torch.float16 if use_fp16 else torch.float32
+
+        self.input_layer = sp.SparseLinear(in_channels, model_channels[0])
+        self.to_latent = sp.SparseLinear(model_channels[-1], 2 * latent_channels)
+        
+        self.blocks = nn.ModuleList([])
+        for i in range(len(num_blocks)):
+            self.blocks.append(nn.ModuleList([]))
+            for j in range(num_blocks[i]):
+                self.blocks[-1].append(
+                    globals()[block_type[i]](
+                        model_channels[i],
+                        **block_args[i],
+                    )
+                )
+            if i < len(num_blocks) - 1:
+                self.blocks[-1].append(
+                    globals()[down_block_type[i]](
+                        model_channels[i],
+                        model_channels[i+1],
+                        **block_args[i],
+                    )
+                )
+                
+        self.initialize_weights()
+        if use_fp16:
+            self.convert_to_fp16()
+
+    @property
+    def device(self) -> torch.device:
+        """
+        Return the device of the model.
+        """
+        return next(self.parameters()).device
+
+    def convert_to_fp16(self) -> None:
+        """
+        Convert the torso of the model to float16.
+        """
+        self.blocks.apply(convert_module_to_f16)
+
+    def convert_to_fp32(self) -> None:
+        """
+        Convert the torso of the model to float32.
+        """
+        self.blocks.apply(convert_module_to_f32)
+
+    def initialize_weights(self) -> None:
+        # Initialize transformer layers:
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+        self.apply(_basic_init)
+
+    def forward(self, x: sp.SparseTensor, sample_posterior=False, return_raw=False):
+        h = self.input_layer(x)
+        h = h.type(self.dtype)
+        for i, res in enumerate(self.blocks):
+            for j, block in enumerate(res):
+                h = block(h)
+        h = h.type(x.dtype)
+        h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
+        h = self.to_latent(h)
+        
+        # Sample from the posterior distribution
+        mean, logvar = h.feats.chunk(2, dim=-1)
+        if sample_posterior:
+            std = torch.exp(0.5 * logvar)
+            z = mean + std * torch.randn_like(std)
+        else:
+            z = mean
+        z = h.replace(z)
+            
+        if return_raw:
+            return z, mean, logvar
+        else:
+            return z
+    
+    
+class SparseUnetVaeDecoder(nn.Module):
+    """
+    Sparse Swin Transformer Unet VAE model.
+    """
+    def __init__(
+        self,
+        out_channels: int,
+        model_channels: List[int],
+        latent_channels: int,
+        num_blocks: List[int],
+        block_type: List[str],
+        up_block_type: List[str],
+        block_args: List[Dict[str, Any]],
+        use_fp16: bool = False,
+        pred_subdiv: bool = True,
+    ):
+        super().__init__()
+        self.out_channels = out_channels
+        self.model_channels = model_channels
+        self.num_blocks = num_blocks
+        self.use_fp16 = use_fp16
+        self.pred_subdiv = pred_subdiv
+        self.dtype = torch.float16 if use_fp16 else torch.float32
+        self.low_vram = False
+        
+        self.output_layer = sp.SparseLinear(model_channels[-1], out_channels)
+        self.from_latent = sp.SparseLinear(latent_channels, model_channels[0])
+        
+        self.blocks = nn.ModuleList([])
+        for i in range(len(num_blocks)):
+            self.blocks.append(nn.ModuleList([]))
+            for j in range(num_blocks[i]):
+                self.blocks[-1].append(
+                    globals()[block_type[i]](
+                        model_channels[i],
+                        **block_args[i],
+                    )
+                )
+            if i < len(num_blocks) - 1:
+                self.blocks[-1].append(
+                    globals()[up_block_type[i]](
+                        model_channels[i],
+                        model_channels[i+1],
+                        pred_subdiv=pred_subdiv,
+                        **block_args[i],
+                    )
+                )
+                    
+        self.initialize_weights()
+        if use_fp16:
+            self.convert_to_fp16()
+            
+    @property
+    def device(self) -> torch.device:
+        """
+        Return the device of the model.
+        """
+        return next(self.parameters()).device
+
+    def convert_to_fp16(self) -> None:
+        """
+        Convert the torso of the model to float16.
+        """
+        self.blocks.apply(convert_module_to_f16)
+
+    def convert_to_fp32(self) -> None:
+        """
+        Convert the torso of the model to float32.
+        """
+        self.blocks.apply(convert_module_to_f32)
+
+    def initialize_weights(self) -> None:
+        # Initialize transformer layers:
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+        self.apply(_basic_init)
+
+    def forward(self, x: sp.SparseTensor, guide_subs: Optional[List[sp.SparseTensor]] = None, return_subs: bool = False) -> sp.SparseTensor:
+        assert guide_subs is None or self.pred_subdiv == False, "Only decoders with pred_subdiv=False can be used with guide_subs"
+        assert return_subs == False or self.pred_subdiv == True, "Only decoders with pred_subdiv=True can be used with return_subs"
+        
+        h = self.from_latent(x)
+        h = h.type(self.dtype)
+        subs_gt = []
+        subs = []
+        for i, res in enumerate(self.blocks):
+            for j, block in enumerate(res):
+                if i < len(self.blocks) - 1 and j == len(res) - 1:
+                    if self.pred_subdiv:
+                        if self.training:
+                            subs_gt.append(h.get_spatial_cache('subdivision'))
+                        h, sub = block(h)
+                        subs.append(sub)
+                    else:
+                        h = block(h, subdiv=guide_subs[i] if guide_subs is not None else None)
+                else:
+                    h = block(h)
+        h = h.type(x.dtype)
+        h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
+        h = self.output_layer(h)
+        if self.training and self.pred_subdiv:
+            return h, subs_gt, subs
+        else:
+            if return_subs:
+                return h, subs
+            else:
+                return h
+    
+    def upsample(self, x: sp.SparseTensor, upsample_times: int) -> torch.Tensor:
+        assert self.pred_subdiv == True, "Only decoders with pred_subdiv=True can be used with upsampling"
+        
+        h = self.from_latent(x)
+        h = h.type(self.dtype)
+        for i, res in enumerate(self.blocks):
+            if i == upsample_times:
+                return h.coords
+            for j, block in enumerate(res):
+                if i < len(self.blocks) - 1 and j == len(res) - 1:
+                    h, sub = block(h)
+                else:
+                    h = block(h)
+       
\ No newline at end of file
diff --git a/pixal3d/models/sparse_elastic_mixin.py b/pixal3d/models/sparse_elastic_mixin.py
new file mode 100644
index 0000000..66d204c
--- /dev/null
+++ b/pixal3d/models/sparse_elastic_mixin.py
@@ -0,0 +1,24 @@
+from contextlib import contextmanager
+from typing import *
+import math
+from ..modules import sparse as sp
+from ..utils.elastic_utils import ElasticModuleMixin
+
+
+class SparseTransformerElasticMixin(ElasticModuleMixin):
+    def _get_input_size(self, x: sp.SparseTensor, *args, **kwargs):
+        return x.feats.shape[0]
+    
+    @contextmanager
+    def with_mem_ratio(self, mem_ratio=1.0):
+        if mem_ratio == 1.0:
+            yield 1.0
+            return
+        num_blocks = len(self.blocks)
+        num_checkpoint_blocks = min(math.ceil((1 - mem_ratio) * num_blocks) + 1, num_blocks)
+        exact_mem_ratio = 1 - (num_checkpoint_blocks - 1) / num_blocks
+        for i in range(num_blocks):
+            self.blocks[i].use_checkpoint = i < num_checkpoint_blocks
+        yield exact_mem_ratio
+        for i in range(num_blocks):
+            self.blocks[i].use_checkpoint = False
diff --git a/pixal3d/models/sparse_structure_flow.py b/pixal3d/models/sparse_structure_flow.py
new file mode 100644
index 0000000..be6b50e
--- /dev/null
+++ b/pixal3d/models/sparse_structure_flow.py
@@ -0,0 +1,298 @@
+from typing import *
+from functools import partial
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from ..modules.utils import convert_module_to, manual_cast, str_to_dtype
+from ..modules.transformer import AbsolutePositionEmbedder, ModulatedTransformerCrossBlock
+from ..modules.attention import RotaryPositionEmbedder
+
+
+class TimestepEmbedder(nn.Module):
+    """
+    Embeds scalar timesteps into vector representations.
+    """
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        """
+        Create sinusoidal timestep embeddings.
+
+        Args:
+            t: a 1-D Tensor of N indices, one per batch element.
+                These may be fractional.
+            dim: the dimension of the output.
+            max_period: controls the minimum frequency of the embeddings.
+
+        Returns:
+            an (N, D) Tensor of positional embeddings.
+        """
+        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+        half = dim // 2
+        freqs = torch.exp(
+            -np.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+        ).to(device=t.device)
+        args = t[:, None].float() * freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+        return embedding
+
+    def forward(self, t):
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        t_emb = self.mlp(t_freq)
+        return t_emb
+
+
+class SparseStructureFlowModel(nn.Module):
+    """
+    Sparse Structure Flow Model for 3D generation.
+    
+    Supports two conditioning modes:
+    - "cross": Standard cross-attention with image features
+    - "proj": View-aligned projection attention with camera-aware features
+    """
+    def __init__(
+        self,
+        resolution: int,
+        in_channels: int,
+        model_channels: int,
+        cond_channels: int,
+        out_channels: int,
+        num_blocks: int,
+        num_heads: Optional[int] = None,
+        num_head_channels: Optional[int] = 64,
+        mlp_ratio: float = 4,
+        pe_mode: Literal["ape", "rope"] = "ape",
+        rope_freq: Tuple[float, float] = (1.0, 10000.0),
+        dtype: str = 'float32',
+        use_checkpoint: bool = False,
+        share_mod: bool = False,
+        initialization: str = 'vanilla',
+        qk_rms_norm: bool = False,
+        qk_rms_norm_cross: bool = False,
+        image_attn_mode: Literal["cross", "proj", "gated_proj"] = "cross",
+        proj_in_channels: Optional[int] = None,
+        vae_in_channels: Optional[int] = None,
+        **kwargs
+    ):
+        super().__init__()
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.cond_channels = cond_channels
+        self.out_channels = out_channels
+        self.num_blocks = num_blocks
+        self.num_heads = num_heads or model_channels // num_head_channels
+        self.mlp_ratio = mlp_ratio
+        self.pe_mode = pe_mode
+        self.use_checkpoint = use_checkpoint
+        self.share_mod = share_mod
+        self.initialization = initialization
+        self.qk_rms_norm = qk_rms_norm
+        self.qk_rms_norm_cross = qk_rms_norm_cross
+        self.image_attn_mode = image_attn_mode
+        self.proj_in_channels = proj_in_channels
+        self.vae_in_channels = vae_in_channels
+        self.dtype = str_to_dtype(dtype)
+
+        self.t_embedder = TimestepEmbedder(model_channels)
+        if share_mod:
+            self.adaLN_modulation = nn.Sequential(
+                nn.SiLU(),
+                nn.Linear(model_channels, 6 * model_channels, bias=True)
+            )
+
+        if pe_mode == "ape":
+            pos_embedder = AbsolutePositionEmbedder(model_channels, 3)
+            coords = torch.meshgrid(*[torch.arange(res, device=self.device) for res in [resolution] * 3], indexing='ij')
+            coords = torch.stack(coords, dim=-1).reshape(-1, 3)
+            pos_emb = pos_embedder(coords)
+            self.register_buffer("pos_emb", pos_emb)
+        elif pe_mode == "rope":
+            pos_embedder = RotaryPositionEmbedder(self.model_channels // self.num_heads, 3)
+            coords = torch.meshgrid(*[torch.arange(res, device=self.device) for res in [resolution] * 3], indexing='ij')
+            coords = torch.stack(coords, dim=-1).reshape(-1, 3)
+            rope_phases = pos_embedder(coords)
+            self.register_buffer("rope_phases", rope_phases)
+            
+        if pe_mode != "rope":
+            self.rope_phases = None
+
+        self.input_layer = nn.Linear(in_channels, model_channels)
+            
+        self.blocks = nn.ModuleList([
+            ModulatedTransformerCrossBlock(
+                model_channels,
+                cond_channels,
+                num_heads=self.num_heads,
+                mlp_ratio=self.mlp_ratio,
+                attn_mode='full',
+                use_checkpoint=self.use_checkpoint,
+                use_rope=(pe_mode == "rope"),
+                rope_freq=rope_freq,
+                share_mod=share_mod,
+                qk_rms_norm=self.qk_rms_norm,
+                qk_rms_norm_cross=self.qk_rms_norm_cross,
+                image_attn_mode=image_attn_mode,
+                proj_in_channels=proj_in_channels,
+                vae_in_channels=vae_in_channels,
+            )
+            for _ in range(num_blocks)
+        ])
+
+        self.out_layer = nn.Linear(model_channels, out_channels)
+
+        self.initialize_weights()
+        self.convert_to(self.dtype)
+
+    @property
+    def device(self) -> torch.device:
+        """
+        Return the device of the model.
+        """
+        return next(self.parameters()).device
+
+    def convert_to(self, dtype: torch.dtype) -> None:
+        """
+        Convert the torso of the model to the specified dtype.
+        """
+        self.dtype = dtype
+        self.blocks.apply(partial(convert_module_to, dtype=dtype))
+
+    def initialize_weights(self) -> None:
+        if self.initialization == 'vanilla':
+            # Initialize transformer layers:
+            def _basic_init(module):
+                if isinstance(module, nn.Linear):
+                    torch.nn.init.xavier_uniform_(module.weight)
+                    if module.bias is not None:
+                        nn.init.constant_(module.bias, 0)
+            self.apply(_basic_init)
+
+            # Initialize timestep embedding MLP:
+            nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+            nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+
+            # Zero-out adaLN modulation layers in DiT blocks:
+            if self.share_mod:
+                nn.init.constant_(self.adaLN_modulation[-1].weight, 0)
+                nn.init.constant_(self.adaLN_modulation[-1].bias, 0)
+            else:
+                for block in self.blocks:
+                    nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+                    nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+
+            # Zero-out output layers:
+            nn.init.constant_(self.out_layer.weight, 0)
+            nn.init.constant_(self.out_layer.bias, 0)
+            
+        elif self.initialization == 'scaled':
+            # Initialize transformer layers:
+            def _basic_init(module):
+                if isinstance(module, nn.Linear):
+                    torch.nn.init.normal_(module.weight, std=np.sqrt(2.0 / (5.0 * self.model_channels)))
+                    if module.bias is not None:
+                        nn.init.constant_(module.bias, 0)
+            self.apply(_basic_init)
+            
+            # Scaled init for to_out and ffn2
+            def _scaled_init(module):
+                if isinstance(module, nn.Linear):
+                    torch.nn.init.normal_(module.weight, std=1.0 / np.sqrt(5 * self.num_blocks * self.model_channels))
+                    if module.bias is not None:
+                        nn.init.constant_(module.bias, 0)
+            for block in self.blocks:
+                block.self_attn.to_out.apply(_scaled_init)
+                # Handle cross, proj, and gated_proj modes
+                if self.image_attn_mode in ("proj", "gated_proj"):
+                    block.cross_attn.cross_attn_block.to_out.apply(_scaled_init)
+                else:
+                    block.cross_attn.to_out.apply(_scaled_init)
+                block.mlp.mlp[2].apply(_scaled_init)
+            
+            # Initialize input layer to make the initial representation have variance 1
+            nn.init.normal_(self.input_layer.weight, std=1.0 / np.sqrt(self.in_channels))
+            nn.init.zeros_(self.input_layer.bias)
+            
+            # Initialize timestep embedding MLP:
+            nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+            nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+            
+            # Zero-out adaLN modulation layers in DiT blocks:
+            if self.share_mod:
+                nn.init.constant_(self.adaLN_modulation[-1].weight, 0)
+                nn.init.constant_(self.adaLN_modulation[-1].bias, 0)
+            else:
+                for block in self.blocks:
+                    nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+                    nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+
+            # Zero-out output layers:
+            nn.init.constant_(self.out_layer.weight, 0)
+            nn.init.constant_(self.out_layer.bias, 0)
+
+    def forward(self, x: torch.Tensor, t: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+        """
+        Forward pass.
+        
+        Args:
+            x: Input tensor [B, C, D, H, W]
+            t: Timestep tensor [B]
+            cond: Conditioning tensor. For "cross" mode: [B, N, D]. 
+                  For "proj" mode: dict {'global': global_cond, 'proj': proj_cond} 
+                  or tuple of (global_cond, proj_cond)
+        
+        Returns:
+            Output tensor [B, C, D, H, W]
+        """
+        assert [*x.shape] == [x.shape[0], self.in_channels, *[self.resolution] * 3], \
+                f"Input shape mismatch, got {x.shape}, expected {[x.shape[0], self.in_channels, *[self.resolution] * 3]}"
+
+        h = x.view(*x.shape[:2], -1).permute(0, 2, 1).contiguous()
+
+        h = self.input_layer(h)
+        if self.pe_mode == "ape":
+            h = h + self.pos_emb[None]
+        t_emb = self.t_embedder(t)
+        if self.share_mod:
+            t_emb = self.adaLN_modulation(t_emb)
+        t_emb = manual_cast(t_emb, self.dtype)
+        h = manual_cast(h, self.dtype)
+        
+        # Handle different conditioning modes
+        if self.image_attn_mode == 'proj':
+            if isinstance(cond, dict):
+                global_cond = cond['global']
+                proj_cond = cond['proj']
+            else:
+                global_cond, proj_cond = cond
+            global_cond = manual_cast(global_cond, self.dtype)
+            proj_cond = manual_cast(proj_cond, self.dtype)
+            cond = (global_cond, proj_cond)
+        elif self.image_attn_mode == 'gated_proj':
+            global_cond = manual_cast(cond['global'], self.dtype)
+            proj_semantic = manual_cast(cond['proj_semantic'], self.dtype)
+            proj_color = manual_cast(cond['proj_color'], self.dtype)
+            cond = {'global': global_cond, 'proj_semantic': proj_semantic, 'proj_color': proj_color}
+        else:
+            cond = manual_cast(cond, self.dtype)
+            
+        for block in self.blocks:
+            h = block(h, t_emb, cond, self.rope_phases)
+        h = manual_cast(h, x.dtype)
+        h = F.layer_norm(h, h.shape[-1:])
+        h = self.out_layer(h)
+
+        h = h.permute(0, 2, 1).view(h.shape[0], h.shape[2], *[self.resolution] * 3).contiguous()
+
+        return h
diff --git a/pixal3d/models/sparse_structure_vae.py b/pixal3d/models/sparse_structure_vae.py
new file mode 100644
index 0000000..c3e0913
--- /dev/null
+++ b/pixal3d/models/sparse_structure_vae.py
@@ -0,0 +1,306 @@
+from typing import *
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from ..modules.norm import GroupNorm32, ChannelLayerNorm32
+from ..modules.spatial import pixel_shuffle_3d
+from ..modules.utils import zero_module, convert_module_to_f16, convert_module_to_f32
+
+
+def norm_layer(norm_type: str, *args, **kwargs) -> nn.Module:
+    """
+    Return a normalization layer.
+    """
+    if norm_type == "group":
+        return GroupNorm32(32, *args, **kwargs)
+    elif norm_type == "layer":
+        return ChannelLayerNorm32(*args, **kwargs)
+    else:
+        raise ValueError(f"Invalid norm type {norm_type}")
+
+
+class ResBlock3d(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        out_channels: Optional[int] = None,
+        norm_type: Literal["group", "layer"] = "layer",
+    ):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+
+        self.norm1 = norm_layer(norm_type, channels)
+        self.norm2 = norm_layer(norm_type, self.out_channels)
+        self.conv1 = nn.Conv3d(channels, self.out_channels, 3, padding=1)
+        self.conv2 = zero_module(nn.Conv3d(self.out_channels, self.out_channels, 3, padding=1))
+        self.skip_connection = nn.Conv3d(channels, self.out_channels, 1) if channels != self.out_channels else nn.Identity()
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = self.norm1(x)
+        h = F.silu(h)
+        h = self.conv1(h)
+        h = self.norm2(h)
+        h = F.silu(h)
+        h = self.conv2(h)
+        h = h + self.skip_connection(x)
+        return h
+
+
+class DownsampleBlock3d(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        mode: Literal["conv", "avgpool"] = "conv",
+    ):
+        assert mode in ["conv", "avgpool"], f"Invalid mode {mode}"
+
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        if mode == "conv":
+            self.conv = nn.Conv3d(in_channels, out_channels, 2, stride=2)
+        elif mode == "avgpool":
+            assert in_channels == out_channels, "Pooling mode requires in_channels to be equal to out_channels"
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if hasattr(self, "conv"):
+            return self.conv(x)
+        else:
+            return F.avg_pool3d(x, 2)
+
+
+class UpsampleBlock3d(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        mode: Literal["conv", "nearest"] = "conv",
+    ):
+        assert mode in ["conv", "nearest"], f"Invalid mode {mode}"
+
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        if mode == "conv":
+            self.conv = nn.Conv3d(in_channels, out_channels*8, 3, padding=1)
+        elif mode == "nearest":
+            assert in_channels == out_channels, "Nearest mode requires in_channels to be equal to out_channels"
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if hasattr(self, "conv"):
+            x = self.conv(x)
+            return pixel_shuffle_3d(x, 2)
+        else:
+            return F.interpolate(x, scale_factor=2, mode="nearest")
+        
+
+class SparseStructureEncoder(nn.Module):
+    """
+    Encoder for Sparse Structure (\mathcal{E}_S in the paper Sec. 3.3).
+    
+    Args:
+        in_channels (int): Channels of the input.
+        latent_channels (int): Channels of the latent representation.
+        num_res_blocks (int): Number of residual blocks at each resolution.
+        channels (List[int]): Channels of the encoder blocks.
+        num_res_blocks_middle (int): Number of residual blocks in the middle.
+        norm_type (Literal["group", "layer"]): Type of normalization layer.
+        use_fp16 (bool): Whether to use FP16.
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        latent_channels: int,
+        num_res_blocks: int,
+        channels: List[int],
+        num_res_blocks_middle: int = 2,
+        norm_type: Literal["group", "layer"] = "layer",
+        use_fp16: bool = False,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.latent_channels = latent_channels
+        self.num_res_blocks = num_res_blocks
+        self.channels = channels
+        self.num_res_blocks_middle = num_res_blocks_middle
+        self.norm_type = norm_type
+        self.use_fp16 = use_fp16
+        self.dtype = torch.float16 if use_fp16 else torch.float32
+
+        self.input_layer = nn.Conv3d(in_channels, channels[0], 3, padding=1)
+
+        self.blocks = nn.ModuleList([])
+        for i, ch in enumerate(channels):
+            self.blocks.extend([
+                ResBlock3d(ch, ch)
+                for _ in range(num_res_blocks)
+            ])
+            if i < len(channels) - 1:
+                self.blocks.append(
+                    DownsampleBlock3d(ch, channels[i+1])
+                )
+        
+        self.middle_block = nn.Sequential(*[
+            ResBlock3d(channels[-1], channels[-1])
+            for _ in range(num_res_blocks_middle)
+        ])
+
+        self.out_layer = nn.Sequential(
+            norm_layer(norm_type, channels[-1]),
+            nn.SiLU(),
+            nn.Conv3d(channels[-1], latent_channels*2, 3, padding=1)
+        )
+
+        if use_fp16:
+            self.convert_to_fp16()
+
+    @property
+    def device(self) -> torch.device:
+        """
+        Return the device of the model.
+        """
+        return next(self.parameters()).device
+
+    def convert_to_fp16(self) -> None:
+        """
+        Convert the torso of the model to float16.
+        """
+        self.use_fp16 = True
+        self.dtype = torch.float16
+        self.blocks.apply(convert_module_to_f16)
+        self.middle_block.apply(convert_module_to_f16)
+
+    def convert_to_fp32(self) -> None:
+        """
+        Convert the torso of the model to float32.
+        """
+        self.use_fp16 = False
+        self.dtype = torch.float32
+        self.blocks.apply(convert_module_to_f32)
+        self.middle_block.apply(convert_module_to_f32)
+
+    def forward(self, x: torch.Tensor, sample_posterior: bool = False, return_raw: bool = False) -> torch.Tensor:
+        h = self.input_layer(x)
+        h = h.type(self.dtype)
+
+        for block in self.blocks:
+            h = block(h)
+        h = self.middle_block(h)
+
+        h = h.type(x.dtype)
+        h = self.out_layer(h)
+
+        mean, logvar = h.chunk(2, dim=1)
+
+        if sample_posterior:
+            std = torch.exp(0.5 * logvar)
+            z = mean + std * torch.randn_like(std)
+        else:
+            z = mean
+            
+        if return_raw:
+            return z, mean, logvar
+        return z
+        
+
+class SparseStructureDecoder(nn.Module):
+    """
+    Decoder for Sparse Structure (\mathcal{D}_S in the paper Sec. 3.3).
+    
+    Args:
+        out_channels (int): Channels of the output.
+        latent_channels (int): Channels of the latent representation.
+        num_res_blocks (int): Number of residual blocks at each resolution.
+        channels (List[int]): Channels of the decoder blocks.
+        num_res_blocks_middle (int): Number of residual blocks in the middle.
+        norm_type (Literal["group", "layer"]): Type of normalization layer.
+        use_fp16 (bool): Whether to use FP16.
+    """ 
+    def __init__(
+        self,
+        out_channels: int,
+        latent_channels: int,
+        num_res_blocks: int,
+        channels: List[int],
+        num_res_blocks_middle: int = 2,
+        norm_type: Literal["group", "layer"] = "layer",
+        use_fp16: bool = False,
+    ):
+        super().__init__()
+        self.out_channels = out_channels
+        self.latent_channels = latent_channels
+        self.num_res_blocks = num_res_blocks
+        self.channels = channels
+        self.num_res_blocks_middle = num_res_blocks_middle
+        self.norm_type = norm_type
+        self.use_fp16 = use_fp16
+        self.dtype = torch.float16 if use_fp16 else torch.float32
+
+        self.input_layer = nn.Conv3d(latent_channels, channels[0], 3, padding=1)
+
+        self.middle_block = nn.Sequential(*[
+            ResBlock3d(channels[0], channels[0])
+            for _ in range(num_res_blocks_middle)
+        ])
+
+        self.blocks = nn.ModuleList([])
+        for i, ch in enumerate(channels):
+            self.blocks.extend([
+                ResBlock3d(ch, ch)
+                for _ in range(num_res_blocks)
+            ])
+            if i < len(channels) - 1:
+                self.blocks.append(
+                    UpsampleBlock3d(ch, channels[i+1])
+                )
+
+        self.out_layer = nn.Sequential(
+            norm_layer(norm_type, channels[-1]),
+            nn.SiLU(),
+            nn.Conv3d(channels[-1], out_channels, 3, padding=1)
+        )
+
+        if use_fp16:
+            self.convert_to_fp16()
+
+    @property
+    def device(self) -> torch.device:
+        """
+        Return the device of the model.
+        """
+        return next(self.parameters()).device
+    
+    def convert_to_fp16(self) -> None:
+        """
+        Convert the torso of the model to float16.
+        """
+        self.use_fp16 = True
+        self.dtype = torch.float16
+        self.blocks.apply(convert_module_to_f16)
+        self.middle_block.apply(convert_module_to_f16)
+
+    def convert_to_fp32(self) -> None:
+        """
+        Convert the torso of the model to float32.
+        """
+        self.use_fp16 = False
+        self.dtype = torch.float32
+        self.blocks.apply(convert_module_to_f32)
+        self.middle_block.apply(convert_module_to_f32)
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = self.input_layer(x)
+        
+        h = h.type(self.dtype)
+                
+        h = self.middle_block(h)
+        for block in self.blocks:
+            h = block(h)
+
+        h = h.type(x.dtype)
+        h = self.out_layer(h)
+        return h
diff --git a/pixal3d/models/structured_latent_flow.py b/pixal3d/models/structured_latent_flow.py
new file mode 100644
index 0000000..23d78ad
--- /dev/null
+++ b/pixal3d/models/structured_latent_flow.py
@@ -0,0 +1,265 @@
+from typing import *
+from functools import partial
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from ..modules.utils import convert_module_to, manual_cast, str_to_dtype
+from ..modules.transformer import AbsolutePositionEmbedder
+from ..modules import sparse as sp
+from ..modules.sparse.transformer import ModulatedSparseTransformerCrossBlock
+from .sparse_structure_flow import TimestepEmbedder
+from .sparse_elastic_mixin import SparseTransformerElasticMixin
+    
+
+class SLatFlowModel(nn.Module):
+    """
+    Structured Latent Flow Model for 3D generation.
+    
+    Supports two conditioning modes:
+    - "cross": Standard cross-attention with image features
+    - "proj": View-aligned projection attention with camera-aware features
+    """
+    def __init__(
+        self,
+        resolution: int,
+        in_channels: int,
+        model_channels: int,
+        cond_channels: int,
+        out_channels: int,
+        num_blocks: int,
+        num_heads: Optional[int] = None,
+        num_head_channels: Optional[int] = 64,
+        mlp_ratio: float = 4,
+        pe_mode: Literal["ape", "rope"] = "ape",
+        rope_freq: Tuple[float, float] = (1.0, 10000.0),
+        dtype: str = 'float32',
+        use_checkpoint: bool = False,
+        share_mod: bool = False,
+        initialization: str = 'vanilla',
+        qk_rms_norm: bool = False,
+        qk_rms_norm_cross: bool = False,
+        image_attn_mode: Literal["cross", "proj", "gated_proj"] = "cross",
+        proj_in_channels: Optional[int] = None,
+        vae_in_channels: Optional[int] = None,
+        **kwargs
+    ):
+        super().__init__()
+        self.resolution = resolution
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.cond_channels = cond_channels
+        self.out_channels = out_channels
+        self.num_blocks = num_blocks
+        self.num_heads = num_heads or model_channels // num_head_channels
+        self.mlp_ratio = mlp_ratio
+        self.pe_mode = pe_mode
+        self.use_checkpoint = use_checkpoint
+        self.share_mod = share_mod
+        self.initialization = initialization
+        self.qk_rms_norm = qk_rms_norm
+        self.qk_rms_norm_cross = qk_rms_norm_cross
+        self.image_attn_mode = image_attn_mode
+        self.proj_in_channels = proj_in_channels
+        self.vae_in_channels = vae_in_channels
+        self.dtype = str_to_dtype(dtype)
+
+        self.t_embedder = TimestepEmbedder(model_channels)
+        if share_mod:
+            self.adaLN_modulation = nn.Sequential(
+                nn.SiLU(),
+                nn.Linear(model_channels, 6 * model_channels, bias=True)
+            )
+
+        if pe_mode == "ape":
+            self.pos_embedder = AbsolutePositionEmbedder(model_channels)
+
+        self.input_layer = sp.SparseLinear(in_channels, model_channels)
+            
+        self.blocks = nn.ModuleList([
+            ModulatedSparseTransformerCrossBlock(
+                model_channels,
+                cond_channels,
+                num_heads=self.num_heads,
+                mlp_ratio=self.mlp_ratio,
+                attn_mode='full',
+                use_checkpoint=self.use_checkpoint,
+                use_rope=(pe_mode == "rope"),
+                rope_freq=rope_freq,
+                share_mod=self.share_mod,
+                qk_rms_norm=self.qk_rms_norm,
+                qk_rms_norm_cross=self.qk_rms_norm_cross,
+                image_attn_mode=image_attn_mode,
+                proj_in_channels=proj_in_channels,
+                vae_in_channels=vae_in_channels,
+            )
+            for _ in range(num_blocks)
+        ])
+            
+        self.out_layer = sp.SparseLinear(model_channels, out_channels)
+
+        self.initialize_weights()
+        self.convert_to(self.dtype)
+
+    @property
+    def device(self) -> torch.device:
+        """
+        Return the device of the model.
+        """
+        return next(self.parameters()).device
+
+    def convert_to(self, dtype: torch.dtype) -> None:
+        """
+        Convert the torso of the model to the specified dtype.
+        """
+        self.dtype = dtype
+        self.blocks.apply(partial(convert_module_to, dtype=dtype))
+
+    def initialize_weights(self) -> None:
+        if self.initialization == 'vanilla':
+            # Initialize transformer layers:
+            def _basic_init(module):
+                if isinstance(module, nn.Linear):
+                    torch.nn.init.xavier_uniform_(module.weight)
+                    if module.bias is not None:
+                        nn.init.constant_(module.bias, 0)
+            self.apply(_basic_init)
+
+            # Initialize timestep embedding MLP:
+            nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+            nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+
+            # Zero-out adaLN modulation layers in DiT blocks:
+            if self.share_mod:
+                nn.init.constant_(self.adaLN_modulation[-1].weight, 0)
+                nn.init.constant_(self.adaLN_modulation[-1].bias, 0)
+            else:
+                for block in self.blocks:
+                    nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+                    nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+
+            # Zero-out output layers:
+            nn.init.constant_(self.out_layer.weight, 0)
+            nn.init.constant_(self.out_layer.bias, 0)
+            
+        elif self.initialization == 'scaled':
+            # Initialize transformer layers:
+            def _basic_init(module):
+                if isinstance(module, nn.Linear):
+                    torch.nn.init.normal_(module.weight, std=np.sqrt(2.0 / (5.0 * self.model_channels)))
+                    if module.bias is not None:
+                        nn.init.constant_(module.bias, 0)
+            self.apply(_basic_init)
+            
+            # Scaled init for to_out and ffn2
+            def _scaled_init(module):
+                if isinstance(module, nn.Linear):
+                    torch.nn.init.normal_(module.weight, std=1.0 / np.sqrt(5 * self.num_blocks * self.model_channels))
+                    if module.bias is not None:
+                        nn.init.constant_(module.bias, 0)
+            for block in self.blocks:
+                block.self_attn.to_out.apply(_scaled_init)
+                # Handle cross, proj, and gated_proj modes
+                if self.image_attn_mode in ("proj", "gated_proj"):
+                    block.cross_attn.cross_attn_block.to_out.apply(_scaled_init)
+                else:
+                    block.cross_attn.to_out.apply(_scaled_init)
+                block.mlp.mlp[2].apply(_scaled_init)
+            
+            # Initialize input layer to make the initial representation have variance 1
+            nn.init.normal_(self.input_layer.weight, std=1.0 / np.sqrt(self.in_channels))
+            nn.init.zeros_(self.input_layer.bias)
+            
+            # Initialize timestep embedding MLP:
+            nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
+            nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
+            
+            # Zero-out adaLN modulation layers in DiT blocks:
+            if self.share_mod:
+                nn.init.constant_(self.adaLN_modulation[-1].weight, 0)
+                nn.init.constant_(self.adaLN_modulation[-1].bias, 0)
+            else:
+                for block in self.blocks:
+                    nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+                    nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+
+            # Zero-out output layers:
+            nn.init.constant_(self.out_layer.weight, 0)
+            nn.init.constant_(self.out_layer.bias, 0)
+
+    def forward(
+        self,
+        x: sp.SparseTensor,
+        t: torch.Tensor,
+        cond: Union[torch.Tensor, List[torch.Tensor], Dict[str, Union[torch.Tensor, sp.SparseTensor]], Tuple],
+        concat_cond: Optional[sp.SparseTensor] = None,
+        **kwargs
+    ) -> sp.SparseTensor:
+        """
+        Forward pass.
+        
+        Args:
+            x: SparseTensor input
+            t: Timestep tensor [B]
+            cond: Conditioning tensor. For "cross" mode: list of tensors or tensor.
+                  For "proj" mode: dict {'global': global_cond, 'proj': proj_cond} 
+                  or tuple of (global_cond, proj_cond)
+            concat_cond: Optional concatenation condition
+        
+        Returns:
+            SparseTensor output
+        """
+        if concat_cond is not None:
+            x = sp.sparse_cat([x, concat_cond], dim=-1)
+
+        h = self.input_layer(x)
+        h = manual_cast(h, self.dtype)
+        t_emb = self.t_embedder(t)
+        if self.share_mod:
+            t_emb = self.adaLN_modulation(t_emb)
+        t_emb = manual_cast(t_emb, self.dtype)
+
+        if self.pe_mode == "ape":
+            pe = self.pos_embedder(h.coords[:, 1:])
+            h = h + manual_cast(pe, self.dtype)
+            
+        # Handle different conditioning modes
+        if self.image_attn_mode == 'proj':
+            if isinstance(cond, dict):
+                global_cond = cond['global']
+                proj_cond = cond['proj']
+            else:
+                global_cond, proj_cond = cond
+            if isinstance(global_cond, list):
+                global_cond = sp.VarLenTensor.from_tensor_list(global_cond)
+            global_cond = manual_cast(global_cond, self.dtype)
+            proj_cond = manual_cast(proj_cond, self.dtype)
+            cond = (global_cond, proj_cond)
+        elif self.image_attn_mode == 'gated_proj':
+            global_cond = cond['global']
+            if isinstance(global_cond, list):
+                global_cond = sp.VarLenTensor.from_tensor_list(global_cond)
+            global_cond = manual_cast(global_cond, self.dtype)
+            proj_semantic = manual_cast(cond['proj_semantic'], self.dtype)
+            proj_color = manual_cast(cond['proj_color'], self.dtype)
+            cond = {'global': global_cond, 'proj_semantic': proj_semantic, 'proj_color': proj_color}
+        else:
+            if isinstance(cond, list):
+                cond = sp.VarLenTensor.from_tensor_list(cond)
+            cond = manual_cast(cond, self.dtype)
+
+        for block in self.blocks:
+            h = block(h, t_emb, cond)
+
+        h = manual_cast(h, x.dtype)
+        h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
+        h = self.out_layer(h)
+        return h
+
+
+class ElasticSLatFlowModel(SparseTransformerElasticMixin, SLatFlowModel):
+    """
+    SLat Flow Model with elastic memory management.
+    Used for training with low VRAM.
+    """
+    pass
diff --git a/pixal3d/modules/attention/__init__.py b/pixal3d/modules/attention/__init__.py
new file mode 100644
index 0000000..318c6d3
--- /dev/null
+++ b/pixal3d/modules/attention/__init__.py
@@ -0,0 +1,4 @@
+from .full_attn import *
+from .modules import *
+from .rope import *
+from .proj_attention import ProjectAttention, GatedProjectAttention
diff --git a/pixal3d/modules/attention/config.py b/pixal3d/modules/attention/config.py
new file mode 100644
index 0000000..579db83
--- /dev/null
+++ b/pixal3d/modules/attention/config.py
@@ -0,0 +1,32 @@
+from typing import *
+
+BACKEND = 'flash_attn' 
+DEBUG = False
+
+def __from_env():
+    import os
+    
+    global BACKEND
+    global DEBUG
+    
+    env_attn_backend = os.environ.get('ATTN_BACKEND')
+    env_attn_debug = os.environ.get('ATTN_DEBUG')
+    
+    if env_attn_backend is not None and env_attn_backend in ['xformers', 'flash_attn', 'flash_attn_3', 'flash_attn_4', 'sdpa', 'naive']:
+        BACKEND = env_attn_backend
+    if env_attn_debug is not None:
+        DEBUG = env_attn_debug == '1'
+
+    print(f"[ATTENTION] Using backend: {BACKEND}")
+        
+
+__from_env()
+    
+
+def set_backend(backend: Literal['xformers', 'flash_attn', 'flash_attn_3', 'flash_attn_4', 'sdpa', 'naive']):
+    global BACKEND
+    BACKEND = backend
+
+def set_debug(debug: bool):
+    global DEBUG
+    DEBUG = debug
diff --git a/pixal3d/modules/attention/full_attn.py b/pixal3d/modules/attention/full_attn.py
new file mode 100644
index 0000000..3c02162
--- /dev/null
+++ b/pixal3d/modules/attention/full_attn.py
@@ -0,0 +1,153 @@
+from typing import *
+import torch
+import math
+from . import config
+
+
+__all__ = [
+    'scaled_dot_product_attention',
+]
+
+
+def _naive_sdpa(q, k, v):
+    """
+    Naive implementation of scaled dot product attention.
+    """
+    q = q.permute(0, 2, 1, 3)   # [N, H, L, C]
+    k = k.permute(0, 2, 1, 3)   # [N, H, L, C]
+    v = v.permute(0, 2, 1, 3)   # [N, H, L, C]
+    scale_factor = 1 / math.sqrt(q.size(-1))
+    attn_weight = q @ k.transpose(-2, -1) * scale_factor
+    attn_weight = torch.softmax(attn_weight, dim=-1)
+    out = attn_weight @ v
+    out = out.permute(0, 2, 1, 3)   # [N, L, H, C]
+    return out
+
+
+@overload
+def scaled_dot_product_attention(qkv: torch.Tensor) -> torch.Tensor:
+    """
+    Apply scaled dot product attention.
+
+    Args:
+        qkv (torch.Tensor): A [N, L, 3, H, C] tensor containing Qs, Ks, and Vs.
+    """
+    ...
+
+@overload
+def scaled_dot_product_attention(q: torch.Tensor, kv: torch.Tensor) -> torch.Tensor:
+    """
+    Apply scaled dot product attention.
+
+    Args:
+        q (torch.Tensor): A [N, L, H, C] tensor containing Qs.
+        kv (torch.Tensor): A [N, L, 2, H, C] tensor containing Ks and Vs.
+    """
+    ...
+
+@overload
+def scaled_dot_product_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+    """
+    Apply scaled dot product attention.
+
+    Args:
+        q (torch.Tensor): A [N, L, H, Ci] tensor containing Qs.
+        k (torch.Tensor): A [N, L, H, Ci] tensor containing Ks.
+        v (torch.Tensor): A [N, L, H, Co] tensor containing Vs.
+
+    Note:
+        k and v are assumed to have the same coordinate map.
+    """
+    ...
+
+def scaled_dot_product_attention(*args, **kwargs):
+    arg_names_dict = {
+        1: ['qkv'],
+        2: ['q', 'kv'],
+        3: ['q', 'k', 'v']
+    }
+    num_all_args = len(args) + len(kwargs)
+    assert num_all_args in arg_names_dict, f"Invalid number of arguments, got {num_all_args}, expected 1, 2, or 3"
+    for key in arg_names_dict[num_all_args][len(args):]:
+        assert key in kwargs, f"Missing argument {key}"
+
+    if num_all_args == 1:
+        qkv = args[0] if len(args) > 0 else kwargs['qkv']
+        assert len(qkv.shape) == 5 and qkv.shape[2] == 3, f"Invalid shape for qkv, got {qkv.shape}, expected [N, L, 3, H, C]"
+        device = qkv.device
+
+    elif num_all_args == 2:
+        q = args[0] if len(args) > 0 else kwargs['q']
+        kv = args[1] if len(args) > 1 else kwargs['kv']
+        assert q.shape[0] == kv.shape[0], f"Batch size mismatch, got {q.shape[0]} and {kv.shape[0]}"
+        assert len(q.shape) == 4, f"Invalid shape for q, got {q.shape}, expected [N, L, H, C]"
+        assert len(kv.shape) == 5, f"Invalid shape for kv, got {kv.shape}, expected [N, L, 2, H, C]"
+        device = q.device
+
+    elif num_all_args == 3:
+        q = args[0] if len(args) > 0 else kwargs['q']
+        k = args[1] if len(args) > 1 else kwargs['k']
+        v = args[2] if len(args) > 2 else kwargs['v']
+        assert q.shape[0] == k.shape[0] == v.shape[0], f"Batch size mismatch, got {q.shape[0]}, {k.shape[0]}, and {v.shape[0]}"
+        assert len(q.shape) == 4, f"Invalid shape for q, got {q.shape}, expected [N, L, H, Ci]"
+        assert len(k.shape) == 4, f"Invalid shape for k, got {k.shape}, expected [N, L, H, Ci]"
+        assert len(v.shape) == 4, f"Invalid shape for v, got {v.shape}, expected [N, L, H, Co]"
+        device = q.device    
+
+    if config.BACKEND == 'xformers':
+        if 'xops' not in globals():
+            import xformers.ops as xops
+        if num_all_args == 1:
+            q, k, v = qkv.unbind(dim=2)
+        elif num_all_args == 2:
+            k, v = kv.unbind(dim=2)
+        out = xops.memory_efficient_attention(q, k, v)
+    elif config.BACKEND == 'flash_attn':
+        if 'flash_attn' not in globals():
+            import flash_attn
+        if num_all_args == 1:
+            out = flash_attn.flash_attn_qkvpacked_func(qkv)
+        elif num_all_args == 2:
+            out = flash_attn.flash_attn_kvpacked_func(q, kv)
+        elif num_all_args == 3:
+            out = flash_attn.flash_attn_func(q, k, v)
+    elif config.BACKEND == 'flash_attn_3':
+        if 'flash_attn_3' not in globals():
+            import flash_attn_interface as flash_attn_3
+            if num_all_args == 1:
+                out = flash_attn_3.flash_attn_qkvpacked_func(qkv)
+            elif num_all_args == 2:
+                k, v = kv.unbind(dim=2)
+                out = flash_attn_3.flash_attn_func(q, k, v)
+            elif num_all_args == 3:
+                out = flash_attn_3.flash_attn_func(q, k, v)
+    elif config.BACKEND == 'flash_attn_4':
+        if 'flash_attn_4' not in globals():
+            from flash_attn.cute import flash_attn_func as flash_attn_4_func
+        if num_all_args == 1:
+            q, k, v = qkv.unbind(dim=2)
+        elif num_all_args == 2:
+            k, v = kv.unbind(dim=2)
+        out, _ = flash_attn_4_func(q, k, v)
+    elif config.BACKEND == 'sdpa':
+        if 'sdpa' not in globals():
+            from torch.nn.functional import scaled_dot_product_attention as sdpa
+        if num_all_args == 1:
+            q, k, v = qkv.unbind(dim=2)
+        elif num_all_args == 2:
+            k, v = kv.unbind(dim=2)
+        q = q.permute(0, 2, 1, 3)   # [N, H, L, C]
+        k = k.permute(0, 2, 1, 3)   # [N, H, L, C]
+        v = v.permute(0, 2, 1, 3)   # [N, H, L, C]
+        out = sdpa(q, k, v)         # [N, H, L, C]
+        out = out.permute(0, 2, 1, 3)   # [N, L, H, C]
+    elif config.BACKEND == 'naive':
+        if num_all_args == 1:
+            q, k, v = qkv.unbind(dim=2)
+        elif num_all_args == 2:
+            k, v = kv.unbind(dim=2)
+        out = _naive_sdpa(q, k, v)
+    else:
+        raise ValueError(f"Unknown attention module: {config.BACKEND}")
+    
+    return out
diff --git a/pixal3d/modules/attention/modules.py b/pixal3d/modules/attention/modules.py
new file mode 100644
index 0000000..492784c
--- /dev/null
+++ b/pixal3d/modules/attention/modules.py
@@ -0,0 +1,102 @@
+from typing import *
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .full_attn import scaled_dot_product_attention
+from .rope import RotaryPositionEmbedder
+
+
+class MultiHeadRMSNorm(nn.Module):
+    def __init__(self, dim: int, heads: int):
+        super().__init__()
+        self.scale = dim ** 0.5
+        self.gamma = nn.Parameter(torch.ones(heads, dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return (F.normalize(x.float(), dim = -1) * self.gamma * self.scale).to(x.dtype)
+    
+
+class MultiHeadAttention(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        num_heads: int,
+        ctx_channels: Optional[int]=None,
+        type: Literal["self", "cross"] = "self",
+        attn_mode: Literal["full", "windowed"] = "full",
+        window_size: Optional[int] = None,
+        shift_window: Optional[Tuple[int, int, int]] = None,
+        qkv_bias: bool = True,
+        use_rope: bool = False,
+        rope_freq: Tuple[float, float] = (1.0, 10000.0),
+        qk_rms_norm: bool = False,
+    ):
+        super().__init__()
+        assert channels % num_heads == 0
+        assert type in ["self", "cross"], f"Invalid attention type: {type}"
+        assert attn_mode in ["full", "windowed"], f"Invalid attention mode: {attn_mode}"
+        assert type == "self" or attn_mode == "full", "Cross-attention only supports full attention"
+        
+        if attn_mode == "windowed":
+            raise NotImplementedError("Windowed attention is not yet implemented")
+        
+        self.channels = channels
+        self.head_dim = channels // num_heads
+        self.ctx_channels = ctx_channels if ctx_channels is not None else channels
+        self.num_heads = num_heads
+        self._type = type
+        self.attn_mode = attn_mode
+        self.window_size = window_size
+        self.shift_window = shift_window
+        self.use_rope = use_rope
+        self.qk_rms_norm = qk_rms_norm
+
+        if self._type == "self":
+            self.to_qkv = nn.Linear(channels, channels * 3, bias=qkv_bias)
+        else:
+            self.to_q = nn.Linear(channels, channels, bias=qkv_bias)
+            self.to_kv = nn.Linear(self.ctx_channels, channels * 2, bias=qkv_bias)
+            
+        if self.qk_rms_norm:
+            self.q_rms_norm = MultiHeadRMSNorm(self.head_dim, num_heads)
+            self.k_rms_norm = MultiHeadRMSNorm(self.head_dim, num_heads)
+            
+        self.to_out = nn.Linear(channels, channels)
+    
+    def forward(self, x: torch.Tensor, context: Optional[torch.Tensor] = None, phases: Optional[torch.Tensor] = None) -> torch.Tensor:
+        B, L, C = x.shape
+        if self._type == "self":
+            qkv = self.to_qkv(x)
+            qkv = qkv.reshape(B, L, 3, self.num_heads, -1)
+            
+            if self.attn_mode == "full":
+                if self.qk_rms_norm or self.use_rope:
+                    q, k, v = qkv.unbind(dim=2)
+                    if self.qk_rms_norm:
+                        q = self.q_rms_norm(q)
+                        k = self.k_rms_norm(k)
+                    if self.use_rope:
+                        assert phases is not None, "Phases must be provided for RoPE"
+                        q = RotaryPositionEmbedder.apply_rotary_embedding(q, phases)
+                        k = RotaryPositionEmbedder.apply_rotary_embedding(k, phases)
+                    h = scaled_dot_product_attention(q, k, v)
+                else:
+                    h = scaled_dot_product_attention(qkv)
+            elif self.attn_mode == "windowed":
+                raise NotImplementedError("Windowed attention is not yet implemented")
+        else:
+            Lkv = context.shape[1]
+            q = self.to_q(x)
+            kv = self.to_kv(context)
+            q = q.reshape(B, L, self.num_heads, -1)
+            kv = kv.reshape(B, Lkv, 2, self.num_heads, -1)
+            if self.qk_rms_norm:
+                q = self.q_rms_norm(q)
+                k, v = kv.unbind(dim=2)
+                k = self.k_rms_norm(k)
+                h = scaled_dot_product_attention(q, k, v)
+            else:
+                h = scaled_dot_product_attention(q, kv)
+        h = h.reshape(B, L, -1)
+        h = self.to_out(h)
+        return h
diff --git a/pixal3d/modules/attention/proj_attention.py b/pixal3d/modules/attention/proj_attention.py
new file mode 100644
index 0000000..d788a98
--- /dev/null
+++ b/pixal3d/modules/attention/proj_attention.py
@@ -0,0 +1,101 @@
+"""
+View-Aligned Projection Attention Module for Pixal3D
+
+This module implements the projection-based attention mechanism that combines
+global cross-attention with view-aligned projected features.
+
+Supports two modes:
+- "proj": Standard projection (DINOv3 only), per-block proj_linear
+- "gated_proj": Gated fusion of DINOv3 (semantic) + VAE (color) features
+"""
+
+from typing import *
+import torch
+import torch.nn as nn
+
+
+class ProjectAttention(nn.Module):
+    """
+    Projection-based Attention Module with per-block proj_linear.
+    
+    Combines global cross-attention with view-aligned projected features.
+    Each block owns a proj_linear that projects DINOv3 features from
+    proj_in_channels (e.g. 1024) to model_channels (e.g. 1536).
+    
+    The module receives:
+    - x: Input features from the transformer
+    - context: A dict with keys:
+      - 'global': Global image features, shape [B, M, ctx_channels]
+      - 'proj': View-aligned projected features, shape [B, N, proj_in_channels]
+    
+    The output combines the cross-attention result with the projected context.
+    """
+    def __init__(self, cross_attn_block: nn.Module, channels: int, proj_in_channels: int):
+        super().__init__()
+        self.cross_attn_block = cross_attn_block
+        self.proj_linear = nn.Linear(proj_in_channels, channels, bias=True)
+        
+    def forward(self, x: torch.Tensor, context: Union[Dict[str, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]) -> torch.Tensor:
+        if isinstance(context, dict):
+            global_context = context['global']
+            proj_context = context['proj']
+        else:
+            global_context, proj_context = context
+        
+        global_out = self.cross_attn_block(x, global_context)
+        proj_out = self.proj_linear(proj_context)
+        context_combined = proj_out + global_out
+        return context_combined
+
+
+class GatedProjectAttention(nn.Module):
+    """
+    Concat-Projection Attention Module for DINOv3 (semantic) + VAE (color) features.
+    
+    Concatenates DINOv3 and VAE projected features and applies a single linear
+    projection to model_channels. This is mathematically equivalent to two
+    separate proj_linears + addition, but allows cross-dimensional interactions
+    between semantic and color features through the shared weight matrix.
+    
+    Zero-initialized for stable training: at init, fused=0 so only global
+    cross-attention contributes; color+semantic signals are gradually learned.
+    
+    The module receives:
+    - x: Input features from the transformer
+    - context: A dict with keys:
+      - 'global': Global image features, shape [B, M, ctx_channels]
+      - 'proj_semantic': DINOv3 projected features, shape [B, N, dino_channels]
+      - 'proj_color': VAE projected features, shape [B, N, vae_channels]
+    """
+    def __init__(
+        self,
+        cross_attn_block: nn.Module,
+        channels: int,
+        dino_in_channels: int,
+        vae_in_channels: int,
+    ):
+        """
+        Args:
+            cross_attn_block: The underlying cross-attention module
+            channels: Model channels (output dimension)
+            dino_in_channels: DINOv3 proj feature dimension (e.g. 1024)
+            vae_in_channels: VAE latent feature dimension (e.g. 16)
+        """
+        super().__init__()
+        self.cross_attn_block = cross_attn_block
+        self.proj_linear = nn.Linear(dino_in_channels + vae_in_channels, channels, bias=True)
+        # Zero-init: at start, fused=0, only global cross-attn contributes
+        nn.init.zeros_(self.proj_linear.weight)
+        nn.init.zeros_(self.proj_linear.bias)
+        
+    def forward(self, x: torch.Tensor, context: Union[Dict[str, torch.Tensor], Tuple]) -> torch.Tensor:
+        if isinstance(context, dict):
+            global_context = context['global']
+            proj_semantic = context['proj_semantic']
+            proj_color = context['proj_color']
+        else:
+            global_context, proj_semantic, proj_color = context
+        
+        global_out = self.cross_attn_block(x, global_context)
+        fused = self.proj_linear(torch.cat([proj_semantic, proj_color], dim=-1))
+        return fused + global_out
diff --git a/pixal3d/modules/attention/rope.py b/pixal3d/modules/attention/rope.py
new file mode 100644
index 0000000..1cf6c5b
--- /dev/null
+++ b/pixal3d/modules/attention/rope.py
@@ -0,0 +1,48 @@
+from typing import *
+import torch
+import torch.nn as nn
+
+
+class RotaryPositionEmbedder(nn.Module):
+    def __init__(
+        self, 
+        head_dim: int,
+        dim: int = 3,
+        rope_freq: Tuple[float, float] = (1.0, 10000.0)
+    ):
+        super().__init__()
+        assert head_dim % 2 == 0, "Head dim must be divisible by 2"
+        self.head_dim = head_dim
+        self.dim = dim
+        self.rope_freq = rope_freq
+        self.freq_dim = head_dim // 2 // dim
+        self.freqs = torch.arange(self.freq_dim, dtype=torch.float32) / self.freq_dim
+        self.freqs = rope_freq[0] / (rope_freq[1] ** (self.freqs))
+        
+    def _get_phases(self, indices: torch.Tensor) -> torch.Tensor:
+        self.freqs = self.freqs.to(indices.device)
+        phases = torch.outer(indices, self.freqs)
+        phases = torch.polar(torch.ones_like(phases), phases)
+        return phases
+    
+    @staticmethod
+    def apply_rotary_embedding(x: torch.Tensor, phases: torch.Tensor) -> torch.Tensor:
+        x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+        x_rotated = x_complex * phases.unsqueeze(-2)
+        x_embed = torch.view_as_real(x_rotated).reshape(*x_rotated.shape[:-1], -1).to(x.dtype)
+        return x_embed
+        
+    def forward(self, indices: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            indices (torch.Tensor): [..., N, C] tensor of spatial positions
+        """
+        assert indices.shape[-1] == self.dim, f"Last dim of indices must be {self.dim}"
+        phases = self._get_phases(indices.reshape(-1)).reshape(*indices.shape[:-1], -1)
+        if phases.shape[-1] < self.head_dim // 2:
+                padn = self.head_dim // 2 - phases.shape[-1]
+                phases = torch.cat([phases, torch.polar(
+                    torch.ones(*phases.shape[:-1], padn, device=phases.device),
+                    torch.zeros(*phases.shape[:-1], padn, device=phases.device)
+                )], dim=-1)
+        return phases
\ No newline at end of file
diff --git a/pixal3d/modules/image_feature_extractor.py b/pixal3d/modules/image_feature_extractor.py
new file mode 100644
index 0000000..c3cb515
--- /dev/null
+++ b/pixal3d/modules/image_feature_extractor.py
@@ -0,0 +1,118 @@
+from typing import *
+import torch
+import torch.nn.functional as F
+from torchvision import transforms
+from transformers import DINOv3ViTModel
+import numpy as np
+from PIL import Image
+
+
+class DinoV2FeatureExtractor:
+    """
+    Feature extractor for DINOv2 models.
+    """
+    def __init__(self, model_name: str):
+        self.model_name = model_name
+        self.model = torch.hub.load('facebookresearch/dinov2', model_name, pretrained=True)
+        self.model.eval()
+        self.transform = transforms.Compose([
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ])
+
+    def to(self, device):
+        self.model.to(device)
+
+    def cuda(self):
+        self.model.cuda()
+
+    def cpu(self):
+        self.model.cpu()
+    
+    @torch.no_grad()
+    def __call__(self, image: Union[torch.Tensor, List[Image.Image]]) -> torch.Tensor:
+        """
+        Extract features from the image.
+        
+        Args:
+            image: A batch of images as a tensor of shape (B, C, H, W) or a list of PIL images.
+        
+        Returns:
+            A tensor of shape (B, N, D) where N is the number of patches and D is the feature dimension.
+        """
+        if isinstance(image, torch.Tensor):
+            assert image.ndim == 4, "Image tensor should be batched (B, C, H, W)"
+        elif isinstance(image, list):
+            assert all(isinstance(i, Image.Image) for i in image), "Image list should be list of PIL images"
+            image = [i.resize((518, 518), Image.LANCZOS) for i in image]
+            image = [np.array(i.convert('RGB')).astype(np.float32) / 255 for i in image]
+            image = [torch.from_numpy(i).permute(2, 0, 1).float() for i in image]
+            image = torch.stack(image).cuda()
+        else:
+            raise ValueError(f"Unsupported type of image: {type(image)}")
+        
+        image = self.transform(image).cuda()
+        features = self.model(image, is_training=True)['x_prenorm']
+        patchtokens = F.layer_norm(features, features.shape[-1:])
+        return patchtokens
+    
+
+class DinoV3FeatureExtractor:
+    """
+    Feature extractor for DINOv3 models.
+    """
+    def __init__(self, model_name: str, image_size=512):
+        self.model_name = model_name
+        self.model = DINOv3ViTModel.from_pretrained(model_name)
+        self.model.eval()
+        self.image_size = image_size
+        self.transform = transforms.Compose([
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ])
+
+    def to(self, device):
+        self.model.to(device)
+
+    def cuda(self):
+        self.model.cuda()
+
+    def cpu(self):
+        self.model.cpu()
+
+    def extract_features(self, image: torch.Tensor) -> torch.Tensor:
+        image = image.to(self.model.embeddings.patch_embeddings.weight.dtype)
+        hidden_states = self.model.embeddings(image, bool_masked_pos=None)
+        position_embeddings = self.model.rope_embeddings(image)
+
+        for i, layer_module in enumerate(self.model.layer):
+            hidden_states = layer_module(
+                hidden_states,
+                position_embeddings=position_embeddings,
+            )
+
+        return F.layer_norm(hidden_states, hidden_states.shape[-1:])
+        
+    @torch.no_grad()
+    def __call__(self, image: Union[torch.Tensor, List[Image.Image]]) -> torch.Tensor:
+        """
+        Extract features from the image.
+        
+        Args:
+            image: A batch of images as a tensor of shape (B, C, H, W) or a list of PIL images.
+        
+        Returns:
+            A tensor of shape (B, N, D) where N is the number of patches and D is the feature dimension.
+        """
+        if isinstance(image, torch.Tensor):
+            assert image.ndim == 4, "Image tensor should be batched (B, C, H, W)"
+        elif isinstance(image, list):
+            assert all(isinstance(i, Image.Image) for i in image), "Image list should be list of PIL images"
+            image = [i.resize((self.image_size, self.image_size), Image.LANCZOS) for i in image]
+            image = [np.array(i.convert('RGB')).astype(np.float32) / 255 for i in image]
+            image = [torch.from_numpy(i).permute(2, 0, 1).float() for i in image]
+            image = torch.stack(image).cuda()
+        else:
+            raise ValueError(f"Unsupported type of image: {type(image)}")
+        
+        image = self.transform(image).cuda()
+        features = self.extract_features(image)
+        return features
diff --git a/pixal3d/modules/norm.py b/pixal3d/modules/norm.py
new file mode 100644
index 0000000..78675d0
--- /dev/null
+++ b/pixal3d/modules/norm.py
@@ -0,0 +1,32 @@
+import torch
+import torch.nn as nn
+from .utils import manual_cast
+
+
+class LayerNorm32(nn.LayerNorm):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x_dtype = x.dtype
+        x = manual_cast(x, torch.float32)
+        o = super().forward(x)
+        return manual_cast(o, x_dtype)
+    
+
+class GroupNorm32(nn.GroupNorm):
+    """
+    A GroupNorm layer that converts to float32 before the forward pass.
+    """
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x_dtype = x.dtype
+        x = manual_cast(x, torch.float32)
+        o = super().forward(x)
+        return manual_cast(o, x_dtype)
+    
+    
+class ChannelLayerNorm32(LayerNorm32):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        DIM = x.dim()
+        x = x.permute(0, *range(2, DIM), 1).contiguous()
+        x = super().forward(x)
+        x = x.permute(0, DIM-1, *range(1, DIM-1)).contiguous()
+        return x
+    
\ No newline at end of file
diff --git a/pixal3d/modules/sparse/__init__.py b/pixal3d/modules/sparse/__init__.py
new file mode 100644
index 0000000..e73f232
--- /dev/null
+++ b/pixal3d/modules/sparse/__init__.py
@@ -0,0 +1,69 @@
+from . import config
+import importlib
+
+__attributes = {
+    'VarLenTensor': 'basic',
+    'varlen_cat': 'basic',
+    'varlen_unbind': 'basic',
+    'SparseTensor': 'basic',
+    'sparse_cat': 'basic',
+    'sparse_unbind': 'basic',
+    'SparseGroupNorm': 'norm',
+    'SparseLayerNorm': 'norm',
+    'SparseGroupNorm32': 'norm',
+    'SparseLayerNorm32': 'norm',
+    'SparseReLU': 'nonlinearity',
+    'SparseSiLU': 'nonlinearity',
+    'SparseGELU': 'nonlinearity',
+    'SparseActivation': 'nonlinearity',
+    'SparseLinear': 'linear',
+    'sparse_scaled_dot_product_attention': 'attention',
+    'SerializeMode': 'attention',
+    'sparse_serialized_scaled_dot_product_self_attention': 'attention',
+    'sparse_windowed_scaled_dot_product_self_attention': 'attention',
+    'sparse_windowed_scaled_dot_product_cross_attention': 'attention',
+    'SparseRotaryPositionEmbedder':  'attention',
+    'SparseMultiHeadAttention': 'attention',
+    'SparseConv3d': 'conv',
+    'SparseInverseConv3d': 'conv',
+    'SparseDownsample': 'spatial',
+    'SparseUpsample': 'spatial',
+    'SparseSubdivide': 'spatial',
+    'SparseSpatial2Channel': 'spatial',
+    'SparseChannel2Spatial': 'spatial',
+    'sparse_nearest_interpolate': 'spatial',
+    'sparse_trilinear_interpolate': 'spatial',
+    'encode_seq': 'serialize',
+    'decode_seq': 'serialize',
+}
+
+__submodules = ['transformer', 'conv']
+
+__all__ = list(__attributes.keys()) + __submodules
+
+def __getattr__(name):
+    if name not in globals():
+        if name in __attributes:
+            module_name = __attributes[name]
+            module = importlib.import_module(f".{module_name}", __name__)
+            globals()[name] = getattr(module, name)
+        elif name in __submodules:
+            module = importlib.import_module(f".{name}", __name__)
+            globals()[name] = module
+        else:
+            raise AttributeError(f"module {__name__} has no attribute {name}")
+    return globals()[name]
+
+
+# For Pylance
+if __name__ == '__main__':
+    from .basic import *
+    from .norm import *
+    from .nonlinearity import *
+    from .linear import *
+    from .attention import *
+    from .conv import *
+    from .spatial import *
+    from .serialize import *
+    import transformer
+    import conv
diff --git a/pixal3d/modules/sparse/attention/__init__.py b/pixal3d/modules/sparse/attention/__init__.py
new file mode 100644
index 0000000..695105b
--- /dev/null
+++ b/pixal3d/modules/sparse/attention/__init__.py
@@ -0,0 +1,4 @@
+from .full_attn import *
+from .windowed_attn import *
+from .modules import *
+from .proj_attention import *
diff --git a/pixal3d/modules/sparse/attention/full_attn.py b/pixal3d/modules/sparse/attention/full_attn.py
new file mode 100644
index 0000000..363c048
--- /dev/null
+++ b/pixal3d/modules/sparse/attention/full_attn.py
@@ -0,0 +1,255 @@
+from typing import *
+import torch
+from .. import VarLenTensor
+from .. import config
+
+
+__all__ = [
+    'sparse_scaled_dot_product_attention',
+]
+
+
+@overload
+def sparse_scaled_dot_product_attention(qkv: VarLenTensor) -> VarLenTensor:
+    """
+    Apply scaled dot product attention to a sparse tensor.
+
+    Args:
+        qkv (VarLenTensor): A [N, *, 3, H, C] sparse tensor containing Qs, Ks, and Vs.
+    """
+    ...
+
+@overload
+def sparse_scaled_dot_product_attention(q: VarLenTensor, kv: Union[VarLenTensor, torch.Tensor]) -> VarLenTensor:
+    """
+    Apply scaled dot product attention to a sparse tensor.
+
+    Args:
+        q (VarLenTensor): A [N, *, H, C] sparse tensor containing Qs.
+        kv (VarLenTensor or torch.Tensor): A [N, *, 2, H, C] sparse tensor or a [N, L, 2, H, C] dense tensor containing Ks and Vs.
+    """
+    ...
+
+@overload
+def sparse_scaled_dot_product_attention(q: torch.Tensor, kv: VarLenTensor) -> torch.Tensor:
+    """
+    Apply scaled dot product attention to a sparse tensor.
+
+    Args:
+        q (torch.Tensor): A [N, L, H, C] dense tensor containing Qs.
+        kv (VarLenTensor): A [N, *, 2, H, C] sparse tensor containing Ks and Vs.
+    """
+    ...
+
+@overload
+def sparse_scaled_dot_product_attention(q: VarLenTensor, k: VarLenTensor, v: VarLenTensor) -> VarLenTensor:
+    """
+    Apply scaled dot product attention to a sparse tensor.
+
+    Args:
+        q (VarLenTensor): A [N, *, H, Ci] sparse tensor containing Qs.
+        k (VarLenTensor): A [N, *, H, Ci] sparse tensor containing Ks.
+        v (VarLenTensor): A [N, *, H, Co] sparse tensor containing Vs.
+
+    Note:
+        k and v are assumed to have the same coordinate map.
+    """
+    ...
+
+@overload
+def sparse_scaled_dot_product_attention(q: VarLenTensor, k: torch.Tensor, v: torch.Tensor) -> VarLenTensor:
+    """
+    Apply scaled dot product attention to a sparse tensor.
+
+    Args:
+        q (VarLenTensor): A [N, *, H, Ci] sparse tensor containing Qs.
+        k (torch.Tensor): A [N, L, H, Ci] dense tensor containing Ks.
+        v (torch.Tensor): A [N, L, H, Co] dense tensor containing Vs.
+    """
+    ...
+
+@overload
+def sparse_scaled_dot_product_attention(q: torch.Tensor, k: VarLenTensor, v: VarLenTensor) -> torch.Tensor:
+    """
+    Apply scaled dot product attention to a sparse tensor.
+
+    Args:
+        q (torch.Tensor): A [N, L, H, Ci] dense tensor containing Qs.
+        k (VarLenTensor): A [N, *, H, Ci] sparse tensor containing Ks.
+        v (VarLenTensor): A [N, *, H, Co] sparse tensor containing Vs.
+    """
+    ...
+
+def sparse_scaled_dot_product_attention(*args, **kwargs):
+    arg_names_dict = {
+        1: ['qkv'],
+        2: ['q', 'kv'],
+        3: ['q', 'k', 'v']
+    }
+    num_all_args = len(args) + len(kwargs)
+    assert num_all_args in arg_names_dict, f"Invalid number of arguments, got {num_all_args}, expected 1, 2, or 3"
+    for key in arg_names_dict[num_all_args][len(args):]:
+        assert key in kwargs, f"Missing argument {key}"
+
+    if num_all_args == 1:
+        qkv = args[0] if len(args) > 0 else kwargs['qkv']
+        assert isinstance(qkv, VarLenTensor), f"qkv must be a VarLenTensor, got {type(qkv)}"
+        assert len(qkv.shape) == 4 and qkv.shape[1] == 3, f"Invalid shape for qkv, got {qkv.shape}, expected [N, *, 3, H, C]"
+        device = qkv.device
+
+        s = qkv
+        q_seqlen = [qkv.layout[i].stop - qkv.layout[i].start for i in range(qkv.shape[0])]
+        kv_seqlen = q_seqlen
+        qkv = qkv.feats     # [T, 3, H, C]
+
+    elif num_all_args == 2:
+        q = args[0] if len(args) > 0 else kwargs['q']
+        kv = args[1] if len(args) > 1 else kwargs['kv']
+        assert isinstance(q, VarLenTensor) and isinstance(kv, (VarLenTensor, torch.Tensor)) or \
+               isinstance(q, torch.Tensor) and isinstance(kv, VarLenTensor), \
+               f"Invalid types, got {type(q)} and {type(kv)}"
+        assert q.shape[0] == kv.shape[0], f"Batch size mismatch, got {q.shape[0]} and {kv.shape[0]}"
+        device = q.device
+
+        if isinstance(q, VarLenTensor):
+            assert len(q.shape) == 3, f"Invalid shape for q, got {q.shape}, expected [N, *, H, C]"
+            s = q
+            q_seqlen = [q.layout[i].stop - q.layout[i].start for i in range(q.shape[0])]
+            q = q.feats     # [T_Q, H, C]
+        else:
+            assert len(q.shape) == 4, f"Invalid shape for q, got {q.shape}, expected [N, L, H, C]"
+            s = None
+            N, L, H, C = q.shape
+            q_seqlen = [L] * N
+            q = q.reshape(N * L, H, C)   # [T_Q, H, C]
+
+        if isinstance(kv, VarLenTensor):
+            assert len(kv.shape) == 4 and kv.shape[1] == 2, f"Invalid shape for kv, got {kv.shape}, expected [N, *, 2, H, C]"
+            kv_seqlen = [kv.layout[i].stop - kv.layout[i].start for i in range(kv.shape[0])]
+            kv = kv.feats     # [T_KV, 2, H, C]
+        else:
+            assert len(kv.shape) == 5, f"Invalid shape for kv, got {kv.shape}, expected [N, L, 2, H, C]"
+            N, L, _, H, C = kv.shape
+            kv_seqlen = [L] * N
+            kv = kv.reshape(N * L, 2, H, C)   # [T_KV, 2, H, C]
+
+    elif num_all_args == 3:
+        q = args[0] if len(args) > 0 else kwargs['q']
+        k = args[1] if len(args) > 1 else kwargs['k']
+        v = args[2] if len(args) > 2 else kwargs['v']
+        assert isinstance(q, VarLenTensor) and isinstance(k, (VarLenTensor, torch.Tensor)) and type(k) == type(v) or \
+               isinstance(q, torch.Tensor) and isinstance(k, VarLenTensor) and isinstance(v, VarLenTensor), \
+               f"Invalid types, got {type(q)}, {type(k)}, and {type(v)}"
+        assert q.shape[0] == k.shape[0] == v.shape[0], f"Batch size mismatch, got {q.shape[0]}, {k.shape[0]}, and {v.shape[0]}"
+        device = q.device
+
+        if isinstance(q, VarLenTensor):
+            assert len(q.shape) == 3, f"Invalid shape for q, got {q.shape}, expected [N, *, H, Ci]"
+            s = q
+            q_seqlen = [q.layout[i].stop - q.layout[i].start for i in range(q.shape[0])]
+            q = q.feats     # [T_Q, H, Ci]
+        else:
+            assert len(q.shape) == 4, f"Invalid shape for q, got {q.shape}, expected [N, L, H, Ci]"
+            s = None
+            N, L, H, CI = q.shape
+            q_seqlen = [L] * N
+            q = q.reshape(N * L, H, CI)  # [T_Q, H, Ci]
+
+        if isinstance(k, VarLenTensor):
+            assert len(k.shape) == 3, f"Invalid shape for k, got {k.shape}, expected [N, *, H, Ci]"
+            assert len(v.shape) == 3, f"Invalid shape for v, got {v.shape}, expected [N, *, H, Co]"
+            kv_seqlen = [k.layout[i].stop - k.layout[i].start for i in range(k.shape[0])]
+            k = k.feats     # [T_KV, H, Ci]
+            v = v.feats     # [T_KV, H, Co]
+        else:
+            assert len(k.shape) == 4, f"Invalid shape for k, got {k.shape}, expected [N, L, H, Ci]"
+            assert len(v.shape) == 4, f"Invalid shape for v, got {v.shape}, expected [N, L, H, Co]"
+            N, L, H, CI, CO = *k.shape, v.shape[-1]
+            kv_seqlen = [L] * N
+            k = k.reshape(N * L, H, CI)     # [T_KV, H, Ci]
+            v = v.reshape(N * L, H, CO)     # [T_KV, H, Co]
+
+    if config.ATTN == 'xformers':
+        if 'xops' not in globals():
+            import xformers.ops as xops
+        if num_all_args == 1:
+            q, k, v = qkv.unbind(dim=1)
+        elif num_all_args == 2:
+            k, v = kv.unbind(dim=1)
+        q = q.unsqueeze(0)
+        k = k.unsqueeze(0)
+        v = v.unsqueeze(0)
+        mask = xops.fmha.BlockDiagonalMask.from_seqlens(q_seqlen, kv_seqlen)
+        out = xops.memory_efficient_attention(q, k, v, mask)[0]
+    elif config.ATTN == 'flash_attn':
+        if 'flash_attn' not in globals():
+            import flash_attn
+        cu_seqlens_q = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(q_seqlen), dim=0)]).int().to(device)
+        if num_all_args in [2, 3]:
+            cu_seqlens_kv = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(kv_seqlen), dim=0)]).int().to(device)
+        if num_all_args == 1:
+            out = flash_attn.flash_attn_varlen_qkvpacked_func(qkv, cu_seqlens_q, max(q_seqlen))
+        elif num_all_args == 2:
+            out = flash_attn.flash_attn_varlen_kvpacked_func(q, kv, cu_seqlens_q, cu_seqlens_kv, max(q_seqlen), max(kv_seqlen))
+        elif num_all_args == 3:
+            out = flash_attn.flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_kv, max(q_seqlen), max(kv_seqlen))
+    elif config.ATTN == 'flash_attn_3':
+        if 'flash_attn_3' not in globals():
+            import flash_attn_interface as flash_attn_3
+        cu_seqlens_q = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(q_seqlen), dim=0)]).int().to(device)
+        if num_all_args == 1:
+            q, k, v = qkv.unbind(dim=1)
+            cu_seqlens_kv = cu_seqlens_q.clone()
+            max_q_seqlen = max_kv_seqlen = max(q_seqlen)
+        elif num_all_args == 2:
+            k, v = kv.unbind(dim=1)
+            cu_seqlens_kv = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(kv_seqlen), dim=0)]).int().to(device)
+            max_q_seqlen = max(q_seqlen)
+            max_kv_seqlen = max(kv_seqlen)
+        elif num_all_args == 3:
+            cu_seqlens_kv = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(kv_seqlen), dim=0)]).int().to(device)
+            max_q_seqlen = max(q_seqlen)
+            max_kv_seqlen = max(kv_seqlen)
+        out = flash_attn_3.flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_kv, max_q_seqlen, max_kv_seqlen)
+    elif config.ATTN == 'flash_attn_4':
+        if 'flash_attn_4' not in globals():
+            from flash_attn.cute import flash_attn_varlen_func as flash_attn_4_varlen_func
+        cu_seqlens_q = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(q_seqlen), dim=0)]).int().to(device)
+        if num_all_args == 1:
+            q, k, v = qkv.unbind(dim=1)
+            cu_seqlens_kv = cu_seqlens_q.clone()
+            max_q_seqlen = max_kv_seqlen = max(q_seqlen)
+        elif num_all_args == 2:
+            k, v = kv.unbind(dim=1)
+            cu_seqlens_kv = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(kv_seqlen), dim=0)]).int().to(device)
+            max_q_seqlen = max(q_seqlen)
+            max_kv_seqlen = max(kv_seqlen)
+        elif num_all_args == 3:
+            cu_seqlens_kv = torch.cat([torch.tensor([0]), torch.cumsum(torch.tensor(kv_seqlen), dim=0)]).int().to(device)
+            max_q_seqlen = max(q_seqlen)
+            max_kv_seqlen = max(kv_seqlen)
+        out, _ = flash_attn_4_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_kv, max_q_seqlen, max_kv_seqlen)
+    elif config.ATTN == 'sdpa':
+        from torch.nn.functional import scaled_dot_product_attention as _sdpa
+        if num_all_args == 1:
+            q, k, v = qkv.unbind(dim=1)   # [T, H, C] each
+        elif num_all_args == 2:
+            k, v = kv.unbind(dim=1)        # [T_KV, H, C] each
+        # process each batch element independently (no varlen kernel needed)
+        q_offs  = [0] + list(torch.cumsum(torch.tensor(q_seqlen),  dim=0).tolist())
+        kv_offs = [0] + list(torch.cumsum(torch.tensor(kv_seqlen), dim=0).tolist())
+        outs = []
+        for i in range(len(q_seqlen)):
+            qi = q[q_offs[i]:q_offs[i+1]].unsqueeze(0).permute(0, 2, 1, 3)    # [1, H, Lq,  C]
+            ki = k[kv_offs[i]:kv_offs[i+1]].unsqueeze(0).permute(0, 2, 1, 3)  # [1, H, Lkv, C]
+            vi = v[kv_offs[i]:kv_offs[i+1]].unsqueeze(0).permute(0, 2, 1, 3)  # [1, H, Lkv, C]
+            out_i = _sdpa(qi, ki, vi)                                            # [1, H, Lq,  C]
+            outs.append(out_i.permute(0, 2, 1, 3).squeeze(0))                   # [Lq, H, C]
+        out = torch.cat(outs, dim=0)                                             # [T_Q, H, C]
+    else:
+        raise ValueError(f"Unknown attention module: {config.ATTN}")
+
+    if s is not None:
+        return s.replace(out)
+    else:
+        return out.reshape(N, L, H, -1)
diff --git a/pixal3d/modules/sparse/attention/modules.py b/pixal3d/modules/sparse/attention/modules.py
new file mode 100644
index 0000000..d762b4b
--- /dev/null
+++ b/pixal3d/modules/sparse/attention/modules.py
@@ -0,0 +1,141 @@
+from typing import *
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .. import VarLenTensor, SparseTensor
+from .full_attn import sparse_scaled_dot_product_attention
+from .windowed_attn import sparse_windowed_scaled_dot_product_self_attention
+from .rope import SparseRotaryPositionEmbedder
+
+
+class SparseMultiHeadRMSNorm(nn.Module):
+    def __init__(self, dim: int, heads: int):
+        super().__init__()
+        self.scale = dim ** 0.5
+        self.gamma = nn.Parameter(torch.ones(heads, dim))
+
+    def forward(self, x: Union[VarLenTensor, torch.Tensor]) -> Union[VarLenTensor, torch.Tensor]:
+        x_type = x.dtype
+        x = x.float()
+        if isinstance(x, VarLenTensor):
+            x = x.replace(F.normalize(x.feats, dim=-1) * self.gamma * self.scale)
+        else:
+            x = F.normalize(x, dim=-1) * self.gamma * self.scale
+        return x.to(x_type)
+
+
+class SparseMultiHeadAttention(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        num_heads: int,
+        ctx_channels: Optional[int] = None,
+        type: Literal["self", "cross"] = "self",
+        attn_mode: Literal["full", "windowed", "double_windowed"] = "full",
+        window_size: Optional[int] = None,
+        shift_window: Optional[Tuple[int, int, int]] = None,
+        qkv_bias: bool = True,
+        use_rope: bool = False,
+        rope_freq: Tuple[int, int] = (1.0, 10000.0),
+        qk_rms_norm: bool = False,
+    ):
+        super().__init__()
+        assert channels % num_heads == 0
+        assert type in ["self", "cross"], f"Invalid attention type: {type}"
+        assert attn_mode in ["full", "windowed", "double_windowed"], f"Invalid attention mode: {attn_mode}"
+        assert type == "self" or attn_mode == "full", "Cross-attention only supports full attention"
+        assert type == "self" or use_rope is False, "Rotary position embeddings only supported for self-attention"
+        if attn_mode == 'double_windowed':
+            assert window_size % 2 == 0, "Window size must be even for double windowed attention"
+            assert num_heads % 2 == 0, "Number of heads must be even for double windowed attention"
+        self.channels = channels
+        self.head_dim = channels // num_heads
+        self.ctx_channels = ctx_channels if ctx_channels is not None else channels
+        self.num_heads = num_heads
+        self._type = type
+        self.attn_mode = attn_mode
+        self.window_size = window_size
+        self.shift_window = shift_window
+        self.use_rope = use_rope
+        self.qk_rms_norm = qk_rms_norm
+
+        if self._type == "self":
+            self.to_qkv = nn.Linear(channels, channels * 3, bias=qkv_bias)
+        else:
+            self.to_q = nn.Linear(channels, channels, bias=qkv_bias)
+            self.to_kv = nn.Linear(self.ctx_channels, channels * 2, bias=qkv_bias)
+        
+        if self.qk_rms_norm:
+            self.q_rms_norm = SparseMultiHeadRMSNorm(self.head_dim, num_heads)
+            self.k_rms_norm = SparseMultiHeadRMSNorm(self.head_dim, num_heads)
+            
+        self.to_out = nn.Linear(channels, channels)
+
+        if use_rope:
+            self.rope = SparseRotaryPositionEmbedder(self.head_dim, rope_freq=rope_freq)
+
+    @staticmethod
+    def _linear(module: nn.Linear, x: Union[VarLenTensor, torch.Tensor]) -> Union[VarLenTensor, torch.Tensor]:
+        if isinstance(x, VarLenTensor):
+            return x.replace(module(x.feats))
+        else:
+            return module(x)
+
+    @staticmethod
+    def _reshape_chs(x: Union[VarLenTensor, torch.Tensor], shape: Tuple[int, ...]) -> Union[VarLenTensor, torch.Tensor]:
+        if isinstance(x, VarLenTensor):
+            return x.reshape(*shape)
+        else:
+            return x.reshape(*x.shape[:2], *shape)
+
+    def _fused_pre(self, x: Union[VarLenTensor, torch.Tensor], num_fused: int) -> Union[VarLenTensor, torch.Tensor]:
+        if isinstance(x, VarLenTensor):
+            x_feats = x.feats.unsqueeze(0)
+        else:
+            x_feats = x
+        x_feats = x_feats.reshape(*x_feats.shape[:2], num_fused, self.num_heads, -1)
+        return x.replace(x_feats.squeeze(0)) if isinstance(x, VarLenTensor) else x_feats
+    
+    def forward(self, x: SparseTensor, context: Optional[Union[VarLenTensor, torch.Tensor]] = None) -> SparseTensor:
+        if self._type == "self":
+            qkv = self._linear(self.to_qkv, x)
+            qkv = self._fused_pre(qkv, num_fused=3)
+            if self.qk_rms_norm or self.use_rope:
+                q, k, v = qkv.unbind(dim=-3)
+                if self.qk_rms_norm:
+                    q = self.q_rms_norm(q)
+                    k = self.k_rms_norm(k)
+                if self.use_rope:
+                    q, k = self.rope(q, k)
+                qkv = qkv.replace(torch.stack([q.feats, k.feats, v.feats], dim=1))
+            if self.attn_mode == "full":
+                h = sparse_scaled_dot_product_attention(qkv)
+            elif self.attn_mode == "windowed":
+                h = sparse_windowed_scaled_dot_product_self_attention(
+                    qkv, self.window_size, shift_window=self.shift_window
+                )
+            elif self.attn_mode == "double_windowed":
+                qkv0 = qkv.replace(qkv.feats[:, :, self.num_heads//2:])
+                qkv1 = qkv.replace(qkv.feats[:, :, :self.num_heads//2])
+                h0 = sparse_windowed_scaled_dot_product_self_attention(
+                    qkv0, self.window_size, shift_window=(0, 0, 0)
+                )
+                h1 = sparse_windowed_scaled_dot_product_self_attention(
+                    qkv1, self.window_size, shift_window=tuple([self.window_size//2] * 3)
+                )
+                h = qkv.replace(torch.cat([h0.feats, h1.feats], dim=1))
+        else:
+            q = self._linear(self.to_q, x)
+            q = self._reshape_chs(q, (self.num_heads, -1))
+            kv = self._linear(self.to_kv, context)
+            kv = self._fused_pre(kv, num_fused=2)
+            if self.qk_rms_norm:
+                q = self.q_rms_norm(q)
+                k, v = kv.unbind(dim=-3)
+                k = self.k_rms_norm(k)
+                h = sparse_scaled_dot_product_attention(q, k, v)
+            else:
+                h = sparse_scaled_dot_product_attention(q, kv)
+        h = self._reshape_chs(h, (-1,))
+        h = self._linear(self.to_out, h)
+        return h
diff --git a/pixal3d/modules/sparse/attention/proj_attention.py b/pixal3d/modules/sparse/attention/proj_attention.py
new file mode 100644
index 0000000..e29a03e
--- /dev/null
+++ b/pixal3d/modules/sparse/attention/proj_attention.py
@@ -0,0 +1,99 @@
+"""
+Sparse View-Aligned Projection Attention Module for Pixal3D
+
+Sparse versions of ProjectAttention and GatedProjectAttention.
+
+Supports two modes:
+- "proj": Standard projection (DINOv3 only)
+- "gated_proj": Gated fusion of DINOv3 (semantic) + VAE (color) features
+"""
+
+from typing import *
+import torch
+import torch.nn as nn
+from ..basic import SparseTensor, VarLenTensor
+
+
+class SparseProjectAttention(nn.Module):
+    """
+    Sparse Projection-based Attention Module with per-block proj_linear.
+    """
+    def __init__(self, cross_attn_block: nn.Module, channels: int, proj_in_channels: int):
+        super().__init__()
+        self.cross_attn_block = cross_attn_block
+        self.proj_linear = nn.Linear(proj_in_channels, channels, bias=True)
+        
+    def forward(
+        self, 
+        x: SparseTensor, 
+        context: Union[Dict[str, Union[torch.Tensor, VarLenTensor, SparseTensor]], 
+                       Tuple[Union[torch.Tensor, VarLenTensor], SparseTensor]]
+    ) -> SparseTensor:
+        if isinstance(context, dict):
+            global_context = context['global']
+            proj_context = context['proj']
+        else:
+            global_context, proj_context = context
+        
+        global_out = self.cross_attn_block(x, global_context)
+        
+        if isinstance(proj_context, SparseTensor):
+            proj_feats = self.proj_linear(proj_context.feats)
+            combined_feats = proj_feats + global_out.feats
+        else:
+            proj_feats = self.proj_linear(proj_context)
+            combined_feats = proj_feats + global_out.feats
+        
+        return global_out.replace(combined_feats)
+
+
+class SparseGatedProjectAttention(nn.Module):
+    """
+    Sparse Concat-Projection Attention Module for DINOv3 + VAE features.
+    
+    Concatenates DINOv3 and VAE projected features and applies a single linear
+    projection to model_channels. Zero-initialized for stable training.
+    
+    Context dict must contain:
+    - 'global': Global image features for cross-attention
+    - 'proj_semantic': DINOv3 projected features (SparseTensor or Tensor)
+    - 'proj_color': VAE projected features (SparseTensor or Tensor)
+    """
+    def __init__(
+        self,
+        cross_attn_block: nn.Module,
+        channels: int,
+        dino_in_channels: int,
+        vae_in_channels: int,
+    ):
+        super().__init__()
+        self.cross_attn_block = cross_attn_block
+        self.proj_linear = nn.Linear(dino_in_channels + vae_in_channels, channels, bias=True)
+        # Zero-init: at start, fused=0, only global cross-attn contributes
+        nn.init.zeros_(self.proj_linear.weight)
+        nn.init.zeros_(self.proj_linear.bias)
+
+    def _get_feats(self, t):
+        return t.feats if isinstance(t, SparseTensor) else t
+
+    def forward(
+        self,
+        x: SparseTensor,
+        context: Union[Dict[str, Union[torch.Tensor, VarLenTensor, SparseTensor]], Tuple],
+    ) -> SparseTensor:
+        if isinstance(context, dict):
+            global_context = context['global']
+            proj_semantic = context['proj_semantic']
+            proj_color = context['proj_color']
+        else:
+            global_context, proj_semantic, proj_color = context
+
+        global_out = self.cross_attn_block(x, global_context)
+
+        fused = self.proj_linear(torch.cat([
+            self._get_feats(proj_semantic),
+            self._get_feats(proj_color),
+        ], dim=-1))
+        combined_feats = fused + global_out.feats
+
+        return global_out.replace(combined_feats)
diff --git a/pixal3d/modules/sparse/attention/rope.py b/pixal3d/modules/sparse/attention/rope.py
new file mode 100644
index 0000000..fb87729
--- /dev/null
+++ b/pixal3d/modules/sparse/attention/rope.py
@@ -0,0 +1,58 @@
+from typing import *
+import torch
+import torch.nn as nn
+from ..basic import SparseTensor
+
+
+class SparseRotaryPositionEmbedder(nn.Module):
+    def __init__(
+        self, 
+        head_dim: int,
+        dim: int = 3,
+        rope_freq: Tuple[float, float] = (1.0, 10000.0)
+    ):
+        super().__init__()
+        assert head_dim % 2 == 0, "Head dim must be divisible by 2"
+        self.head_dim = head_dim
+        self.dim = dim
+        self.rope_freq = rope_freq
+        self.freq_dim = head_dim // 2 // dim
+        self.freqs = torch.arange(self.freq_dim, dtype=torch.float32) / self.freq_dim
+        self.freqs = rope_freq[0] / (rope_freq[1] ** (self.freqs))
+        
+    def _get_phases(self, indices: torch.Tensor) -> torch.Tensor:
+        self.freqs = self.freqs.to(indices.device)
+        phases = torch.outer(indices, self.freqs)
+        phases = torch.polar(torch.ones_like(phases), phases)
+        return phases
+        
+    def _rotary_embedding(self, x: torch.Tensor, phases: torch.Tensor) -> torch.Tensor:
+        x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+        x_rotated = x_complex * phases.unsqueeze(-2)
+        x_embed = torch.view_as_real(x_rotated).reshape(*x_rotated.shape[:-1], -1).to(x.dtype)
+        return x_embed
+        
+    def forward(self, q: SparseTensor, k: Optional[SparseTensor] = None) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            q (SparseTensor): [..., N, H, D] tensor of queries
+            k (SparseTensor): [..., N, H, D] tensor of keys
+        """
+        assert q.coords.shape[-1] == self.dim + 1, "Last dimension of coords must be equal to dim+1"
+        phases_cache_name = f'rope_phase_{self.dim}d_freq{self.rope_freq[0]}-{self.rope_freq[1]}_hd{self.head_dim}'
+        phases = q.get_spatial_cache(phases_cache_name)
+        if phases is None:
+            coords = q.coords[..., 1:]
+            phases = self._get_phases(coords.reshape(-1)).reshape(*coords.shape[:-1], -1)
+            if phases.shape[-1] < self.head_dim // 2:
+                padn = self.head_dim // 2 - phases.shape[-1]
+                phases = torch.cat([phases, torch.polar(
+                    torch.ones(*phases.shape[:-1], padn, device=phases.device),
+                    torch.zeros(*phases.shape[:-1], padn, device=phases.device)
+                )], dim=-1)
+            q.register_spatial_cache(phases_cache_name, phases)
+        q_embed = q.replace(self._rotary_embedding(q.feats, phases))
+        if k is None:
+            return q_embed
+        k_embed = k.replace(self._rotary_embedding(k.feats, phases))
+        return q_embed, k_embed
\ No newline at end of file
diff --git a/pixal3d/modules/sparse/attention/windowed_attn.py b/pixal3d/modules/sparse/attention/windowed_attn.py
new file mode 100644
index 0000000..675aa08
--- /dev/null
+++ b/pixal3d/modules/sparse/attention/windowed_attn.py
@@ -0,0 +1,205 @@
+from typing import *
+import torch
+import math
+from .. import SparseTensor
+from .. import config
+
+
+__all__ = [
+    'sparse_windowed_scaled_dot_product_self_attention',
+    'sparse_windowed_scaled_dot_product_cross_attention',
+]
+
+
+def calc_window_partition(
+    tensor: SparseTensor,
+    window_size: Union[int, Tuple[int, ...]],
+    shift_window: Union[int, Tuple[int, ...]] = 0,
+) -> Tuple[torch.Tensor, torch.Tensor, List[int], List[int]]:
+    """
+    Calculate serialization and partitioning for a set of coordinates.
+
+    Args:
+        tensor (SparseTensor): The input tensor.
+        window_size (int): The window size to use.
+        shift_window (Tuple[int, ...]): The shift of serialized coordinates.
+
+    Returns:
+        (torch.Tensor): Forwards indices.
+        (torch.Tensor): Backwards indices.
+        (torch.Tensor): Sequence lengths.
+        (dict): Attn func args.
+    """
+    DIM = tensor.coords.shape[1] - 1
+    shift_window = (shift_window,) * DIM if isinstance(shift_window, int) else shift_window
+    window_size = (window_size,) * DIM if isinstance(window_size, int) else window_size
+    shifted_coords = tensor.coords.clone().detach()
+    shifted_coords[:, 1:] += torch.tensor(shift_window, device=tensor.device, dtype=torch.int32).unsqueeze(0)
+
+    MAX_COORDS = [i + j for i, j in zip(tensor.spatial_shape, shift_window)]
+    NUM_WINDOWS = [math.ceil((mc + 1) / ws) for mc, ws in zip(MAX_COORDS, window_size)]
+    OFFSET = torch.cumprod(torch.tensor([1] + NUM_WINDOWS[::-1]), dim=0).tolist()[::-1]
+
+    shifted_coords[:, 1:] //= torch.tensor(window_size, device=tensor.device, dtype=torch.int32).unsqueeze(0)
+    shifted_indices = (shifted_coords * torch.tensor(OFFSET, device=tensor.device, dtype=torch.int32).unsqueeze(0)).sum(dim=1)
+    fwd_indices = torch.argsort(shifted_indices)
+    bwd_indices = torch.empty_like(fwd_indices)
+    bwd_indices[fwd_indices] = torch.arange(fwd_indices.shape[0], device=tensor.device)
+    seq_lens = torch.bincount(shifted_indices)
+    mask = seq_lens != 0
+    seq_lens = seq_lens[mask]
+    
+    if config.ATTN == 'xformers':
+        if 'xops' not in globals():
+            import xformers.ops as xops
+        attn_func_args = {
+            'attn_bias': xops.fmha.BlockDiagonalMask.from_seqlens(seq_lens)
+        }
+    elif config.ATTN == 'flash_attn':
+        attn_func_args = {
+            'cu_seqlens': torch.cat([torch.tensor([0], device=tensor.device), torch.cumsum(seq_lens, dim=0)], dim=0).int(),
+            'max_seqlen': torch.max(seq_lens)
+        }
+    elif config.ATTN == 'flash_attn_4':
+        attn_func_args = {
+            'cu_seqlens_q': torch.cat([torch.tensor([0], device=tensor.device), torch.cumsum(seq_lens, dim=0)], dim=0).int(),
+            'cu_seqlens_k': torch.cat([torch.tensor([0], device=tensor.device), torch.cumsum(seq_lens, dim=0)], dim=0).int(),
+            'max_seqlen_q': torch.max(seq_lens),
+            'max_seqlen_k': torch.max(seq_lens),
+        }
+
+    return fwd_indices, bwd_indices, seq_lens, attn_func_args
+    
+
+def sparse_windowed_scaled_dot_product_self_attention(
+    qkv: SparseTensor,
+    window_size: int,
+    shift_window: Tuple[int, int, int] = (0, 0, 0)
+) -> SparseTensor:
+    """
+    Apply windowed scaled dot product self attention to a sparse tensor.
+
+    Args:
+        qkv (SparseTensor): [N, *, 3, H, C] sparse tensor containing Qs, Ks, and Vs.
+        window_size (int): The window size to use.
+        shift_window (Tuple[int, int, int]): The shift of serialized coordinates.
+        
+    Returns:
+        (SparseTensor): [N, *, H, C] sparse tensor containing the output features.
+    """
+    assert len(qkv.shape) == 4 and qkv.shape[1] == 3, f"Invalid shape for qkv, got {qkv.shape}, expected [N, *, 3, H, C]"
+
+    serialization_spatial_cache_name = f'windowed_attention_{window_size}_{shift_window}'
+    serialization_spatial_cache = qkv.get_spatial_cache(serialization_spatial_cache_name)
+    if serialization_spatial_cache is None:
+        fwd_indices, bwd_indices, seq_lens, attn_func_args = calc_window_partition(qkv, window_size, shift_window)
+        qkv.register_spatial_cache(serialization_spatial_cache_name, (fwd_indices, bwd_indices, seq_lens, attn_func_args))
+    else:
+        fwd_indices, bwd_indices, seq_lens, attn_func_args = serialization_spatial_cache
+    
+    qkv_feats = qkv.feats[fwd_indices]      # [M, 3, H, C]
+
+    if config.DEBUG:
+        start = 0
+        qkv_coords = qkv.coords[fwd_indices]
+        for i in range(len(seq_lens)):
+            seq_coords = qkv_coords[start:start+seq_lens[i]]
+            assert (seq_coords[:, 1:].max(dim=0).values - seq_coords[:, 1:].min(dim=0).values < window_size).all(), \
+                    f"SparseWindowedScaledDotProductSelfAttention: window size exceeded"
+            start += seq_lens[i]
+
+    if config.ATTN == 'xformers':
+        if 'xops' not in globals():
+            import xformers.ops as xops
+        q, k, v = qkv_feats.unbind(dim=1)                                               # [M, H, C]
+        q = q.unsqueeze(0)                                                              # [1, M, H, C]
+        k = k.unsqueeze(0)                                                              # [1, M, H, C]
+        v = v.unsqueeze(0)                                                              # [1, M, H, C]
+        out = xops.memory_efficient_attention(q, k, v, **attn_func_args)[0]             # [M, H, C]
+    elif config.ATTN == 'flash_attn':
+        if 'flash_attn' not in globals():
+            import flash_attn
+        out = flash_attn.flash_attn_varlen_qkvpacked_func(qkv_feats, **attn_func_args)  # [M, H, C]
+
+    out = out[bwd_indices]      # [T, H, C]
+
+    if config.DEBUG:
+        qkv_coords = qkv_coords[bwd_indices]
+        assert torch.equal(qkv_coords, qkv.coords), "SparseWindowedScaledDotProductSelfAttention: coordinate mismatch"
+
+    return qkv.replace(out)
+
+
+def sparse_windowed_scaled_dot_product_cross_attention(
+    q: SparseTensor,
+    kv: SparseTensor,
+    q_window_size: int,
+    kv_window_size: int,
+    q_shift_window: Tuple[int, int, int] = (0, 0, 0),
+    kv_shift_window: Tuple[int, int, int] = (0, 0, 0),
+) -> SparseTensor:
+    """
+    Apply windowed scaled dot product cross attention to two sparse tensors.
+
+    Args:
+        q (SparseTensor): [N, *, H, C] sparse tensor containing Qs.
+        kv (SparseTensor): [N, *, 2, H, C] sparse tensor containing Ks and Vs.
+        q_window_size (int): The window size to use for Qs.
+        kv_window_size (int): The window size to use for Ks and Vs.
+        q_shift_window (Tuple[int, int, int]): The shift of serialized coordinates for Qs.
+        kv_shift_window (Tuple[int, int, int]): The shift of serialized coordinates for Ks and Vs.
+        
+    Returns:
+        (SparseTensor): [N, *, H, C] sparse tensor containing the output features.
+    """
+    assert len(q.shape) == 3, f"Invalid shape for q, got {q.shape}, expected [N, *, H, C]"
+    assert len(kv.shape) == 4 and kv.shape[1] == 2, f"Invalid shape for kv, got {kv.shape}, expected [N, *, 2, H, C]"
+
+    q_serialization_spatial_cache_name = f'windowed_attention_{q_window_size}_{q_shift_window}'
+    q_serialization_spatial_cache = q.get_spatial_cache(q_serialization_spatial_cache_name)
+    if q_serialization_spatial_cache is None:
+        q_fwd_indices, q_bwd_indices, q_seq_lens, q_attn_func_args = calc_window_partition(q, q_window_size, q_shift_window)
+        q.register_spatial_cache(q_serialization_spatial_cache_name, (q_fwd_indices, q_bwd_indices, q_seq_lens, q_attn_func_args))
+    else:
+        q_fwd_indices, q_bwd_indices, q_seq_lens, q_attn_func_args = q_serialization_spatial_cache
+    kv_serialization_spatial_cache_name = f'windowed_attention_{kv_window_size}_{kv_shift_window}'
+    kv_serialization_spatial_cache = kv.get_spatial_cache(kv_serialization_spatial_cache_name)
+    if kv_serialization_spatial_cache is None:
+        kv_fwd_indices, kv_bwd_indices, kv_seq_lens, kv_attn_func_args = calc_window_partition(kv, kv_window_size, kv_shift_window)
+        kv.register_spatial_cache(kv_serialization_spatial_cache_name, (kv_fwd_indices, kv_bwd_indices, kv_seq_lens, kv_attn_func_args))
+    else:
+        kv_fwd_indices, kv_bwd_indices, kv_seq_lens, kv_attn_func_args = kv_serialization_spatial_cache
+
+    assert len(q_seq_lens) == len(kv_seq_lens), "Number of sequences in q and kv must match"
+
+    q_feats = q.feats[q_fwd_indices]      # [M, H, C]
+    kv_feats = kv.feats[kv_fwd_indices]    # [M, 2, H, C]
+
+    if config.ATTN == 'xformers':
+        if 'xops' not in globals():
+            import xformers.ops as xops
+        k, v = kv_feats.unbind(dim=1)                                                   # [M, H, C]
+        q = q.unsqueeze(0)                                                              # [1, M, H, C]
+        k = k.unsqueeze(0)                                                              # [1, M, H, C]
+        v = v.unsqueeze(0)                                                              # [1, M, H, C]
+        mask = xops.fmha.BlockDiagonalMask.from_seqlens(q_seq_lens, kv_seq_lens)
+        out = xops.memory_efficient_attention(q, k, v, attn_bias=mask)[0]               # [M, H, C]
+    elif config.ATTN == 'flash_attn':
+        if 'flash_attn' not in globals():
+            import flash_attn
+        out = flash_attn.flash_attn_varlen_kvpacked_func(q_feats, kv_feats,
+            cu_seqlens_q=q_attn_func_args['cu_seqlens'], cu_seqlens_k=kv_attn_func_args['cu_seqlens'],
+            max_seqlen_q=q_attn_func_args['max_seqlen'], max_seqlen_k=kv_attn_func_args['max_seqlen'],
+        )  # [M, H, C]
+    elif config.ATTN == 'flash_attn_4':
+        if 'flash_attn_4' not in globals():
+            from flash_attn.cute import flash_attn_varlen_func as flash_attn_4_varlen_func
+        k, v = kv_feats.unbind(dim=1)
+        out, _ = flash_attn_4_varlen_func(q_feats, k, v,
+            cu_seqlens_q=q_attn_func_args['cu_seqlens_q'], cu_seqlens_k=kv_attn_func_args['cu_seqlens_k'],
+            max_seqlen_q=q_attn_func_args['max_seqlen_q'], max_seqlen_k=kv_attn_func_args['max_seqlen_k'],
+        )  # [M, H, C]
+
+    out = out[q_bwd_indices]      # [T, H, C]
+
+    return q.replace(out)
diff --git a/pixal3d/modules/sparse/basic.py b/pixal3d/modules/sparse/basic.py
new file mode 100644
index 0000000..880973b
--- /dev/null
+++ b/pixal3d/modules/sparse/basic.py
@@ -0,0 +1,836 @@
+from typing import *
+from fractions import Fraction
+import torch
+from . import config
+
+
+__all__ = [
+    'VarLenTensor',
+    'varlen_cat',
+    'varlen_unbind',
+    'SparseTensor',
+    'sparse_cat',
+    'sparse_unbind',
+]
+
+
+class VarLenTensor:
+    """
+    Sequential tensor with variable length.
+    
+    Args:
+        feats (torch.Tensor): Features of the varlen tensor.
+        layout (List[slice]): Layout of the varlen tensor for each batch
+    """
+    def __init__(self, feats: torch.Tensor, layout: List[slice]=None):
+        self.feats = feats
+        self.layout = layout if layout is not None else [slice(0, feats.shape[0])]
+        self._cache = {}
+        
+    @staticmethod
+    def layout_from_seqlen(seqlen: list) -> List[slice]:
+        """
+        Create a layout from a tensor of sequence lengths.
+        """
+        layout = []
+        start = 0
+        for l in seqlen:
+            layout.append(slice(start, start + l))
+            start += l
+        return layout
+        
+    @staticmethod
+    def from_tensor_list(tensor_list: List[torch.Tensor]) -> 'VarLenTensor':
+        """
+        Create a VarLenTensor from a list of tensors.
+        """
+        feats = torch.cat(tensor_list, dim=0)
+        layout = []
+        start = 0
+        for tensor in tensor_list:
+            layout.append(slice(start, start + tensor.shape[0]))
+            start += tensor.shape[0]
+        return VarLenTensor(feats, layout)
+    
+    def to_tensor_list(self) -> List[torch.Tensor]:
+        """
+        Convert a VarLenTensor to a list of tensors.
+        """
+        tensor_list = []
+        for s in self.layout:
+            tensor_list.append(self.feats[s])
+        return tensor_list
+    
+    def __len__(self) -> int:
+        return len(self.layout)
+    
+    @property
+    def shape(self) -> torch.Size:
+        return torch.Size([len(self.layout), *self.feats.shape[1:]])
+    
+    def dim(self) -> int:
+        return len(self.shape)
+    
+    @property
+    def ndim(self) -> int:
+        return self.dim()
+
+    @property
+    def dtype(self):
+        return self.feats.dtype
+
+    @property
+    def device(self):
+        return self.feats.device
+    
+    @property
+    def seqlen(self) -> torch.LongTensor:
+        if 'seqlen' not in self._cache:
+            self._cache['seqlen'] = torch.tensor([l.stop - l.start for l in self.layout], dtype=torch.long, device=self.device)
+        return self._cache['seqlen']
+    
+    @property
+    def cum_seqlen(self) -> torch.LongTensor:
+        if 'cum_seqlen' not in self._cache:
+            self._cache['cum_seqlen'] = torch.cat([
+                torch.tensor([0], dtype=torch.long, device=self.device),
+                self.seqlen.cumsum(dim=0)
+            ], dim=0)
+        return self._cache['cum_seqlen']
+    
+    @property
+    def batch_boardcast_map(self) -> torch.LongTensor:
+        """
+        Get the broadcast map for the varlen tensor.
+        """
+        if 'batch_boardcast_map' not in self._cache:
+            self._cache['batch_boardcast_map'] = torch.repeat_interleave(
+                torch.arange(len(self.layout), device=self.device),
+                self.seqlen,
+            )
+        return self._cache['batch_boardcast_map']
+    
+    @overload
+    def to(self, dtype: torch.dtype, *, non_blocking: bool = False, copy: bool = False) -> 'VarLenTensor': ...
+
+    @overload
+    def to(self, device: Optional[Union[str, torch.device]] = None, dtype: Optional[torch.dtype] = None, *, non_blocking: bool = False, copy: bool = False) -> 'VarLenTensor': ...
+
+    def to(self, *args, **kwargs) -> 'VarLenTensor':
+        device = None
+        dtype = None
+        if len(args) == 2:
+            device, dtype = args
+        elif len(args) == 1:
+            if isinstance(args[0], torch.dtype):
+                dtype = args[0]
+            else:
+                device = args[0]
+        if 'dtype' in kwargs:
+            assert dtype is None, "to() received multiple values for argument 'dtype'"
+            dtype = kwargs['dtype']
+        if 'device' in kwargs:
+            assert device is None, "to() received multiple values for argument 'device'"
+            device = kwargs['device']
+        non_blocking = kwargs.get('non_blocking', False)
+        copy = kwargs.get('copy', False)
+        
+        new_feats = self.feats.to(device=device, dtype=dtype, non_blocking=non_blocking, copy=copy)
+        return self.replace(new_feats)
+
+    def type(self, dtype):
+        new_feats = self.feats.type(dtype)
+        return self.replace(new_feats)
+
+    def cpu(self) -> 'VarLenTensor':
+        new_feats = self.feats.cpu()
+        return self.replace(new_feats)
+    
+    def cuda(self) -> 'VarLenTensor':
+        new_feats = self.feats.cuda()
+        return self.replace(new_feats)
+
+    def half(self) -> 'VarLenTensor':
+        new_feats = self.feats.half()
+        return self.replace(new_feats)
+    
+    def float(self) -> 'VarLenTensor':
+        new_feats = self.feats.float()
+        return self.replace(new_feats)
+    
+    def detach(self) -> 'VarLenTensor':
+        new_feats = self.feats.detach()
+        return self.replace(new_feats)
+
+    def reshape(self, *shape) -> 'VarLenTensor':
+        new_feats = self.feats.reshape(self.feats.shape[0], *shape)
+        return self.replace(new_feats)
+    
+    def unbind(self, dim: int) -> List['VarLenTensor']:
+        return varlen_unbind(self, dim)
+
+    def replace(self, feats: torch.Tensor) -> 'VarLenTensor':
+        new_tensor = VarLenTensor(
+            feats=feats,
+            layout=self.layout,
+        )
+        new_tensor._cache = self._cache
+        return new_tensor
+    
+    def to_dense(self, max_length=None) -> torch.Tensor:
+        """
+        Convert a VarLenTensor to a dense representation without for-loop.
+        
+        Returns:
+            dense (torch.Tensor): (N, L, C) dense tensor
+            mask (torch.BoolTensor): (N, L) mask indicating valid positions
+        """
+        N = len(self)
+        L = max_length or self.seqlen.max().item()
+        spatial = self.feats.shape[1:]
+        idx = torch.arange(L, device=self.device).unsqueeze(0).expand(N, L)
+        mask = (idx < self.seqlen.unsqueeze(1))
+        mapping = mask.reshape(-1).cumsum(dim=0) - 1
+        dense = self.feats[mapping]
+        dense = dense.reshape(N, L, *spatial)
+        return dense, mask
+
+    def __neg__(self) -> 'VarLenTensor':
+        return self.replace(-self.feats)
+    
+    def __elemwise__(self, other: Union[torch.Tensor, 'VarLenTensor'], op: callable) -> 'VarLenTensor':
+        if isinstance(other, torch.Tensor):
+            try:
+                other = torch.broadcast_to(other, self.shape)
+                other = other[self.batch_boardcast_map]
+            except:
+                pass
+        if isinstance(other, VarLenTensor):
+            other = other.feats
+        new_feats = op(self.feats, other)
+        new_tensor = self.replace(new_feats)
+        return new_tensor
+
+    def __add__(self, other: Union[torch.Tensor, 'VarLenTensor', float]) -> 'VarLenTensor':
+        return self.__elemwise__(other, torch.add)
+
+    def __radd__(self, other: Union[torch.Tensor, 'VarLenTensor', float]) -> 'VarLenTensor':
+        return self.__elemwise__(other, torch.add)
+    
+    def __sub__(self, other: Union[torch.Tensor, 'VarLenTensor', float]) -> 'VarLenTensor':
+        return self.__elemwise__(other, torch.sub)
+    
+    def __rsub__(self, other: Union[torch.Tensor, 'VarLenTensor', float]) -> 'VarLenTensor':
+        return self.__elemwise__(other, lambda x, y: torch.sub(y, x))
+
+    def __mul__(self, other: Union[torch.Tensor, 'VarLenTensor', float]) -> 'VarLenTensor':
+        return self.__elemwise__(other, torch.mul)
+
+    def __rmul__(self, other: Union[torch.Tensor, 'VarLenTensor', float]) -> 'VarLenTensor':
+        return self.__elemwise__(other, torch.mul)
+
+    def __truediv__(self, other: Union[torch.Tensor, 'VarLenTensor', float]) -> 'VarLenTensor':
+        return self.__elemwise__(other, torch.div)
+
+    def __rtruediv__(self, other: Union[torch.Tensor, 'VarLenTensor', float]) -> 'VarLenTensor':
+        return self.__elemwise__(other, lambda x, y: torch.div(y, x))
+
+    def __getitem__(self, idx):
+        if isinstance(idx, int):
+            idx = [idx]
+        elif isinstance(idx, slice):
+            idx = range(*idx.indices(self.shape[0]))
+        elif isinstance(idx, list):
+            assert all(isinstance(i, int) for i in idx), f"Only integer indices are supported: {idx}"
+        elif isinstance(idx, torch.Tensor):
+            if idx.dtype == torch.bool:
+                assert idx.shape == (self.shape[0],), f"Invalid index shape: {idx.shape}"
+                idx = idx.nonzero().squeeze(1)
+            elif idx.dtype in [torch.int32, torch.int64]:
+                assert len(idx.shape) == 1, f"Invalid index shape: {idx.shape}"
+            else:
+                raise ValueError(f"Unknown index type: {idx.dtype}")
+        else:
+            raise ValueError(f"Unknown index type: {type(idx)}")
+        
+        new_feats = []
+        new_layout = []
+        start = 0
+        for new_idx, old_idx in enumerate(idx):
+            new_feats.append(self.feats[self.layout[old_idx]])
+            new_layout.append(slice(start, start + len(new_feats[-1])))
+            start += len(new_feats[-1])
+        new_feats = torch.cat(new_feats, dim=0).contiguous()
+        new_tensor = VarLenTensor(feats=new_feats, layout=new_layout)
+        return new_tensor
+    
+    def reduce(self, op: str, dim: Optional[Union[int, Tuple[int,...]]] = None, keepdim: bool = False) -> torch.Tensor:
+        if isinstance(dim, int):
+            dim = (dim,)
+        
+        if op =='mean':
+            red = self.feats.mean(dim=dim, keepdim=keepdim)
+        elif op =='sum':
+            red = self.feats.sum(dim=dim, keepdim=keepdim)
+        elif op == 'prod':
+            red = self.feats.prod(dim=dim, keepdim=keepdim)
+        else:
+            raise ValueError(f"Unsupported reduce operation: {op}")
+        
+        if dim is None or 0 in dim:
+            return red
+        
+        red = torch.segment_reduce(red, reduce=op, lengths=self.seqlen)
+        return red
+    
+    def mean(self, dim: Optional[Union[int, Tuple[int,...]]] = None, keepdim: bool = False) -> torch.Tensor:
+        return self.reduce(op='mean', dim=dim, keepdim=keepdim)
+        
+    def sum(self, dim: Optional[Union[int, Tuple[int,...]]] = None, keepdim: bool = False) -> torch.Tensor:
+        return self.reduce(op='sum', dim=dim, keepdim=keepdim)
+        
+    def prod(self, dim: Optional[Union[int, Tuple[int,...]]] = None, keepdim: bool = False) -> torch.Tensor:
+        return self.reduce(op='prod', dim=dim, keepdim=keepdim)
+    
+    def std(self, dim: Optional[Union[int, Tuple[int,...]]] = None, keepdim: bool = False) -> torch.Tensor:
+        mean = self.mean(dim=dim, keepdim=True)
+        mean2 = self.replace(self.feats ** 2).mean(dim=dim, keepdim=True)
+        std = (mean2 - mean ** 2).sqrt()
+        return std
+    
+    def __repr__(self) -> str:
+        return f"VarLenTensor(shape={self.shape}, dtype={self.dtype}, device={self.device})"
+
+
+def varlen_cat(inputs: List[VarLenTensor], dim: int = 0) -> VarLenTensor:
+    """
+    Concatenate a list of varlen tensors.
+    
+    Args:
+        inputs (List[VarLenTensor]): List of varlen tensors to concatenate.
+    """
+    if dim == 0:
+        new_feats = torch.cat([input.feats for input in inputs], dim=0)
+        start = 0
+        new_layout = []
+        for input in inputs:
+            for l in input.layout:
+                new_layout.append(slice(start, start + l.stop - l.start))
+                start += l.stop - l.start
+        output = VarLenTensor(feats=new_feats, layout=new_layout)
+    else:
+        feats = torch.cat([input.feats for input in inputs], dim=dim)
+        output = inputs[0].replace(feats)
+
+    return output
+
+
+def varlen_unbind(input: VarLenTensor, dim: int) -> Union[List[VarLenTensor]]:
+    """
+    Unbind a varlen tensor along a dimension.
+    
+    Args:
+        input (VarLenTensor): Varlen tensor to unbind.
+        dim (int): Dimension to unbind.
+    """
+    if dim == 0:
+        return [input[i] for i in range(len(input))]
+    else:
+        feats = input.feats.unbind(dim)
+        return [input.replace(f) for f in feats]
+    
+
+class SparseTensor(VarLenTensor):
+    """
+    Sparse tensor with support for both torchsparse and spconv backends.
+    
+    Parameters:
+    - feats (torch.Tensor): Features of the sparse tensor.
+    - coords (torch.Tensor): Coordinates of the sparse tensor.
+    - shape (torch.Size): Shape of the sparse tensor.
+    - layout (List[slice]): Layout of the sparse tensor for each batch
+    - data (SparseTensorData): Sparse tensor data used for convolusion
+
+    NOTE:
+    - Data corresponding to a same batch should be contiguous.
+    - Coords should be in [0, 1023]
+    """
+    SparseTensorData = None
+
+    @overload
+    def __init__(self, feats: torch.Tensor, coords: torch.Tensor, shape: Optional[torch.Size] = None, **kwargs): ...
+
+    @overload
+    def __init__(self, data, shape: Optional[torch.Size] = None, **kwargs): ...
+
+    def __init__(self, *args, **kwargs):
+        # Lazy import of sparse tensor backend
+        if self.SparseTensorData is None:
+            import importlib
+            if config.CONV == 'torchsparse':
+                self.SparseTensorData = importlib.import_module('torchsparse').SparseTensor
+            elif config.CONV == 'spconv':
+                self.SparseTensorData = importlib.import_module('spconv.pytorch').SparseConvTensor
+                
+        method_id = 0
+        if len(args) != 0:
+            method_id = 0 if isinstance(args[0], torch.Tensor) else 1
+        else:
+            method_id = 1 if 'data' in kwargs else 0
+
+        if method_id == 0:
+            feats, coords, shape = args + (None,) * (3 - len(args))
+            if 'feats' in kwargs:
+                feats = kwargs['feats']
+                del kwargs['feats']
+            if 'coords' in kwargs:
+                coords = kwargs['coords']
+                del kwargs['coords']
+            if 'shape' in kwargs:
+                shape = kwargs['shape']
+                del kwargs['shape']
+
+            if config.CONV == 'torchsparse':
+                self.data = self.SparseTensorData(feats, coords, **kwargs)
+            elif config.CONV == 'spconv':
+                spatial_shape = list(coords.max(0)[0] + 1)
+                self.data = self.SparseTensorData(feats.reshape(feats.shape[0], -1), coords, spatial_shape[1:], spatial_shape[0], **kwargs)
+                self.data._features = feats
+            else:
+                self.data = {
+                    'feats': feats,
+                    'coords': coords,
+                }
+        elif method_id == 1:
+            data, shape = args + (None,) * (2 - len(args))
+            if 'data' in kwargs:
+                data = kwargs['data']
+                del kwargs['data']
+            if 'shape' in kwargs:
+                shape = kwargs['shape']
+                del kwargs['shape']
+
+            self.data = data
+
+        self._shape = shape
+        self._scale = kwargs.get('scale', (Fraction(1, 1), Fraction(1, 1), Fraction(1, 1)))
+        self._spatial_cache = kwargs.get('spatial_cache', {})
+
+        if config.DEBUG:
+            try:
+                assert self.feats.shape[0] == self.coords.shape[0], f"Invalid feats shape: {self.feats.shape}, coords shape: {self.coords.shape}"
+                assert self.shape == self.__cal_shape(self.feats, self.coords), f"Invalid shape: {self.shape}"
+                assert self.layout == self.__cal_layout(self.coords, self.shape[0]), f"Invalid layout: {self.layout}"
+                for i in range(self.shape[0]):
+                    assert torch.all(self.coords[self.layout[i], 0] == i), f"The data of batch {i} is not contiguous"
+            except Exception as e:
+                print('Debugging information:')
+                print(f"- Shape: {self.shape}")
+                print(f"- Layout: {self.layout}")
+                print(f"- Scale: {self._scale}")
+                print(f"- Coords: {self.coords}")
+                raise e
+        
+    @staticmethod
+    def from_tensor_list(feats_list: List[torch.Tensor], coords_list: List[torch.Tensor]) -> 'SparseTensor':
+        """
+        Create a SparseTensor from a list of tensors.
+        """
+        feats = torch.cat(feats_list, dim=0)
+        coords = []
+        for i, coord in enumerate(coords_list):
+            coord = torch.cat([torch.full_like(coord[:, :1], i), coord[:, 1:]], dim=1)
+            coords.append(coord)
+        coords = torch.cat(coords, dim=0)
+        return SparseTensor(feats, coords)
+    
+    def to_tensor_list(self) -> Tuple[List[torch.Tensor], List[torch.Tensor]]:
+        """
+        Convert a SparseTensor to list of tensors.
+        """
+        feats_list = []
+        coords_list = []
+        for s in self.layout:
+            feats_list.append(self.feats[s])
+            coords_list.append(self.coords[s])
+        return feats_list, coords_list
+    
+    def __len__(self) -> int:
+        return len(self.layout)
+        
+    def __cal_shape(self, feats, coords):
+        shape = []
+        shape.append(coords[:, 0].max().item() + 1)
+        shape.extend([*feats.shape[1:]])
+        return torch.Size(shape)
+    
+    def __cal_layout(self, coords, batch_size):
+        seq_len = torch.bincount(coords[:, 0], minlength=batch_size)
+        offset = torch.cumsum(seq_len, dim=0) 
+        layout = [slice((offset[i] - seq_len[i]).item(), offset[i].item()) for i in range(batch_size)]
+        return layout
+    
+    def __cal_spatial_shape(self, coords):
+        return torch.Size((coords[:, 1:].max(0)[0] + 1).tolist())
+    
+    @property
+    def shape(self) -> torch.Size:
+        if self._shape is None:
+            self._shape = self.__cal_shape(self.feats, self.coords)
+        return self._shape
+    
+    @property
+    def layout(self) -> List[slice]:
+        layout = self.get_spatial_cache('layout')
+        if layout is None:
+            layout = self.__cal_layout(self.coords, self.shape[0])
+            self.register_spatial_cache('layout', layout)
+        return layout
+    
+    @property
+    def spatial_shape(self) -> torch.Size:
+        spatial_shape = self.get_spatial_cache('shape')
+        if spatial_shape is None:
+            spatial_shape = self.__cal_spatial_shape(self.coords)
+            self.register_spatial_cache('shape', spatial_shape)
+        return spatial_shape
+
+    @property
+    def feats(self) -> torch.Tensor:
+        if config.CONV == 'torchsparse':
+            return self.data.F
+        elif config.CONV == 'spconv':
+            return self.data.features
+        else:
+            return self.data['feats']
+    
+    @feats.setter
+    def feats(self, value: torch.Tensor):
+        if config.CONV == 'torchsparse':
+            self.data.F = value
+        elif config.CONV == 'spconv':
+            self.data.features = value
+        else:
+            self.data['feats'] = value
+
+    @property
+    def coords(self) -> torch.Tensor:
+        if config.CONV == 'torchsparse':
+            return self.data.C
+        elif config.CONV == 'spconv':
+            return self.data.indices
+        else:
+            return self.data['coords']
+        
+    @coords.setter
+    def coords(self, value: torch.Tensor):
+        if config.CONV == 'torchsparse':
+            self.data.C = value
+        elif config.CONV == 'spconv':
+            self.data.indices = value
+        else:
+            self.data['coords'] = value
+
+    @property
+    def dtype(self):
+        return self.feats.dtype
+
+    @property
+    def device(self):
+        return self.feats.device
+    
+    @property
+    def seqlen(self) -> torch.LongTensor:
+        seqlen = self.get_spatial_cache('seqlen')
+        if seqlen is None:
+            seqlen = torch.tensor([l.stop - l.start for l in self.layout], dtype=torch.long, device=self.device)
+            self.register_spatial_cache('seqlen', seqlen)
+        return seqlen
+    
+    @property
+    def cum_seqlen(self) -> torch.LongTensor:
+        cum_seqlen = self.get_spatial_cache('cum_seqlen')
+        if cum_seqlen is None:
+            cum_seqlen = torch.cat([
+                torch.tensor([0], dtype=torch.long, device=self.device),
+                self.seqlen.cumsum(dim=0)
+            ], dim=0)
+            self.register_spatial_cache('cum_seqlen', cum_seqlen)
+        return cum_seqlen
+    
+    @property
+    def batch_boardcast_map(self) -> torch.LongTensor:
+        """
+        Get the broadcast map for the varlen tensor.
+        """
+        batch_boardcast_map = self.get_spatial_cache('batch_boardcast_map')
+        if batch_boardcast_map is None:
+            batch_boardcast_map = torch.repeat_interleave(
+                torch.arange(len(self.layout), device=self.device),
+                self.seqlen,
+            )
+            self.register_spatial_cache('batch_boardcast_map', batch_boardcast_map)
+        return batch_boardcast_map
+
+    @overload
+    def to(self, dtype: torch.dtype, *, non_blocking: bool = False, copy: bool = False) -> 'SparseTensor': ...
+
+    @overload
+    def to(self, device: Optional[Union[str, torch.device]] = None, dtype: Optional[torch.dtype] = None, *, non_blocking: bool = False, copy: bool = False) -> 'SparseTensor': ...
+
+    def to(self, *args, **kwargs) -> 'SparseTensor':
+        device = None
+        dtype = None
+        if len(args) == 2:
+            device, dtype = args
+        elif len(args) == 1:
+            if isinstance(args[0], torch.dtype):
+                dtype = args[0]
+            else:
+                device = args[0]
+        if 'dtype' in kwargs:
+            assert dtype is None, "to() received multiple values for argument 'dtype'"
+            dtype = kwargs['dtype']
+        if 'device' in kwargs:
+            assert device is None, "to() received multiple values for argument 'device'"
+            device = kwargs['device']
+        non_blocking = kwargs.get('non_blocking', False)
+        copy = kwargs.get('copy', False)
+        
+        new_feats = self.feats.to(device=device, dtype=dtype, non_blocking=non_blocking, copy=copy)
+        new_coords = self.coords.to(device=device, non_blocking=non_blocking, copy=copy)
+        return self.replace(new_feats, new_coords)
+
+    def type(self, dtype):
+        new_feats = self.feats.type(dtype)
+        return self.replace(new_feats)
+
+    def cpu(self) -> 'SparseTensor':
+        new_feats = self.feats.cpu()
+        new_coords = self.coords.cpu()
+        return self.replace(new_feats, new_coords)
+    
+    def cuda(self) -> 'SparseTensor':
+        new_feats = self.feats.cuda()
+        new_coords = self.coords.cuda()
+        return self.replace(new_feats, new_coords)
+
+    def half(self) -> 'SparseTensor':
+        new_feats = self.feats.half()
+        return self.replace(new_feats)
+    
+    def float(self) -> 'SparseTensor':
+        new_feats = self.feats.float()
+        return self.replace(new_feats)
+    
+    def detach(self) -> 'SparseTensor':
+        new_coords = self.coords.detach()
+        new_feats = self.feats.detach()
+        return self.replace(new_feats, new_coords)
+
+    def reshape(self, *shape) -> 'SparseTensor':
+        new_feats = self.feats.reshape(self.feats.shape[0], *shape)
+        return self.replace(new_feats)
+    
+    def unbind(self, dim: int) -> List['SparseTensor']:
+        return sparse_unbind(self, dim)
+
+    def replace(self, feats: torch.Tensor, coords: Optional[torch.Tensor] = None) -> 'SparseTensor':
+        if config.CONV == 'torchsparse':
+            new_data = self.SparseTensorData(
+                feats=feats,
+                coords=self.data.coords if coords is None else coords,
+                stride=self.data.stride,
+                spatial_range=self.data.spatial_range,
+            )
+            new_data._caches = self.data._caches
+        elif config.CONV == 'spconv':
+            new_data = self.SparseTensorData(
+                self.data.features.reshape(self.data.features.shape[0], -1),
+                self.data.indices,
+                self.data.spatial_shape,
+                self.data.batch_size,
+                self.data.grid,
+                self.data.voxel_num,
+                self.data.indice_dict
+            )
+            new_data._features = feats
+            new_data.benchmark = self.data.benchmark
+            new_data.benchmark_record = self.data.benchmark_record
+            new_data.thrust_allocator = self.data.thrust_allocator
+            new_data._timer = self.data._timer
+            new_data.force_algo = self.data.force_algo
+            new_data.int8_scale = self.data.int8_scale
+            if coords is not None:
+                new_data.indices = coords
+        else:
+            new_data = {
+                'feats': feats,
+                'coords': self.data['coords'] if coords is None else coords,
+            }
+        new_tensor = SparseTensor(
+            new_data,
+            shape=torch.Size([self._shape[0]] + list(feats.shape[1:])) if self._shape is not None else None,
+            scale=self._scale,
+            spatial_cache=self._spatial_cache
+        )
+        return new_tensor
+    
+    def to_dense(self) -> torch.Tensor:
+        if config.CONV == 'torchsparse':
+            return self.data.dense()
+        elif config.CONV == 'spconv':
+            return self.data.dense()
+        else:
+            spatial_shape = self.spatial_shape
+            ret = torch.zeros(*self.shape, *spatial_shape, dtype=self.dtype, device=self.device)
+            idx = [self.coords[:, 0], slice(None)] + self.coords[:, 1:].unbind(1)
+            ret[tuple(idx)] = self.feats
+            return ret
+
+    @staticmethod
+    def full(aabb, dim, value, dtype=torch.float32, device=None) -> 'SparseTensor':
+        N, C = dim
+        x = torch.arange(aabb[0], aabb[3] + 1)
+        y = torch.arange(aabb[1], aabb[4] + 1)
+        z = torch.arange(aabb[2], aabb[5] + 1)
+        coords = torch.stack(torch.meshgrid(x, y, z, indexing='ij'), dim=-1).reshape(-1, 3)
+        coords = torch.cat([
+            torch.arange(N).view(-1, 1).repeat(1, coords.shape[0]).view(-1, 1),
+            coords.repeat(N, 1),
+        ], dim=1).to(dtype=torch.int32, device=device)
+        feats = torch.full((coords.shape[0], C), value, dtype=dtype, device=device)
+        return SparseTensor(feats=feats, coords=coords)
+
+    def __merge_sparse_cache(self, other: 'SparseTensor') -> dict:
+        new_cache = {}
+        for k in set(list(self._spatial_cache.keys()) + list(other._spatial_cache.keys())):
+            if k in self._spatial_cache:
+                new_cache[k] = self._spatial_cache[k]
+            if k in other._spatial_cache:
+                if k not in new_cache:
+                    new_cache[k] = other._spatial_cache[k]
+                else:
+                    new_cache[k].update(other._spatial_cache[k])
+        return new_cache
+    
+    def __elemwise__(self, other: Union[torch.Tensor, VarLenTensor], op: callable) -> 'SparseTensor':
+        if isinstance(other, torch.Tensor):
+            try:
+                other = torch.broadcast_to(other, self.shape)
+                other = other[self.batch_boardcast_map]
+            except:
+                pass
+        if isinstance(other, VarLenTensor):
+            other = other.feats
+        new_feats = op(self.feats, other)
+        new_tensor = self.replace(new_feats)
+        if isinstance(other, SparseTensor):
+            new_tensor._spatial_cache = self.__merge_sparse_cache(other)
+        return new_tensor
+
+    def __getitem__(self, idx):
+        if isinstance(idx, int):
+            idx = [idx]
+        elif isinstance(idx, slice):
+            idx = range(*idx.indices(self.shape[0]))
+        elif isinstance(idx, list):
+            assert all(isinstance(i, int) for i in idx), f"Only integer indices are supported: {idx}"
+        elif isinstance(idx, torch.Tensor):
+            if idx.dtype == torch.bool:
+                assert idx.shape == (self.shape[0],), f"Invalid index shape: {idx.shape}"
+                idx = idx.nonzero().squeeze(1)
+            elif idx.dtype in [torch.int32, torch.int64]:
+                assert len(idx.shape) == 1, f"Invalid index shape: {idx.shape}"
+            else:
+                raise ValueError(f"Unknown index type: {idx.dtype}")
+        else:
+            raise ValueError(f"Unknown index type: {type(idx)}")
+        
+        new_coords = []
+        new_feats = []
+        new_layout = []
+        new_shape = torch.Size([len(idx)] + list(self.shape[1:]))
+        start = 0
+        for new_idx, old_idx in enumerate(idx):
+            new_coords.append(self.coords[self.layout[old_idx]].clone())
+            new_coords[-1][:, 0] = new_idx
+            new_feats.append(self.feats[self.layout[old_idx]])
+            new_layout.append(slice(start, start + len(new_coords[-1])))
+            start += len(new_coords[-1])
+        new_coords = torch.cat(new_coords, dim=0).contiguous()
+        new_feats = torch.cat(new_feats, dim=0).contiguous()
+        new_tensor = SparseTensor(feats=new_feats, coords=new_coords, shape=new_shape)
+        new_tensor.register_spatial_cache('layout', new_layout)
+        return new_tensor
+    
+    def clear_spatial_cache(self) -> None:
+        """
+        Clear all spatial caches.
+        """
+        self._spatial_cache = {}
+
+    def register_spatial_cache(self, key, value) -> None:
+        """
+        Register a spatial cache.
+        The spatial cache can be any thing you want to cache.
+        The registery and retrieval of the cache is based on current scale.
+        """
+        scale_key = str(self._scale)
+        if scale_key not in self._spatial_cache:
+            self._spatial_cache[scale_key] = {}
+        self._spatial_cache[scale_key][key] = value
+
+    def get_spatial_cache(self, key=None):
+        """
+        Get a spatial cache.
+        """
+        scale_key = str(self._scale)
+        cur_scale_cache = self._spatial_cache.get(scale_key, {})
+        if key is None:
+            return cur_scale_cache
+        return cur_scale_cache.get(key, None)
+    
+    def __repr__(self) -> str:
+        return f"SparseTensor(shape={self.shape}, dtype={self.dtype}, device={self.device})"
+
+def sparse_cat(inputs: List[SparseTensor], dim: int = 0) -> SparseTensor:
+    """
+    Concatenate a list of sparse tensors.
+    
+    Args:
+        inputs (List[SparseTensor]): List of sparse tensors to concatenate.
+    """
+    if dim == 0:
+        start = 0
+        coords = []
+        for input in inputs:
+            coords.append(input.coords.clone())
+            coords[-1][:, 0] += start
+            start += input.shape[0]
+        coords = torch.cat(coords, dim=0)
+        feats = torch.cat([input.feats for input in inputs], dim=0)
+        output = SparseTensor(
+            coords=coords,
+            feats=feats,
+        )
+    else:
+        feats = torch.cat([input.feats for input in inputs], dim=dim)
+        output = inputs[0].replace(feats)
+
+    return output
+
+
+def sparse_unbind(input: SparseTensor, dim: int) -> List[SparseTensor]:
+    """
+    Unbind a sparse tensor along a dimension.
+    
+    Args:
+        input (SparseTensor): Sparse tensor to unbind.
+        dim (int): Dimension to unbind.
+    """
+    if dim == 0:
+        return [input[i] for i in range(input.shape[0])]
+    else:
+        feats = input.feats.unbind(dim)
+        return [input.replace(f) for f in feats]
diff --git a/pixal3d/modules/sparse/config.py b/pixal3d/modules/sparse/config.py
new file mode 100644
index 0000000..25610e5
--- /dev/null
+++ b/pixal3d/modules/sparse/config.py
@@ -0,0 +1,43 @@
+from typing import *
+
+CONV = 'flex_gemm' 
+DEBUG = False
+ATTN = 'flash_attn'
+
+def __from_env():
+    import os
+    
+    global CONV
+    global DEBUG
+    global ATTN
+    
+    env_sparse_conv_backend = os.environ.get('SPARSE_CONV_BACKEND')
+    env_sparse_debug = os.environ.get('SPARSE_DEBUG')
+    env_sparse_attn_backend = os.environ.get('SPARSE_ATTN_BACKEND')
+    if env_sparse_attn_backend is None:
+        env_sparse_attn_backend = os.environ.get('ATTN_BACKEND')
+
+    if env_sparse_conv_backend is not None and env_sparse_conv_backend in ['none', 'spconv', 'torchsparse', 'flex_gemm']:
+        CONV = env_sparse_conv_backend
+    if env_sparse_debug is not None:
+        DEBUG = env_sparse_debug == '1'
+    if env_sparse_attn_backend is not None and env_sparse_attn_backend in ['xformers', 'flash_attn', 'flash_attn_3', 'flash_attn_4', 'sdpa']:
+        ATTN = env_sparse_attn_backend
+        
+    print(f"[SPARSE] Conv backend: {CONV}; Attention backend: {ATTN}")
+        
+
+__from_env()
+    
+
+def set_conv_backend(backend: Literal['none', 'spconv', 'torchsparse', 'flex_gemm']):
+    global CONV
+    CONV = backend
+
+def set_debug(debug: bool):
+    global DEBUG
+    DEBUG = debug
+
+def set_attn_backend(backend: Literal['xformers', 'flash_attn', 'flash_attn_3', 'flash_attn_4', 'sdpa']):
+    global ATTN
+    ATTN = backend
diff --git a/pixal3d/modules/sparse/conv/__init__.py b/pixal3d/modules/sparse/conv/__init__.py
new file mode 100644
index 0000000..a7f5911
--- /dev/null
+++ b/pixal3d/modules/sparse/conv/__init__.py
@@ -0,0 +1,2 @@
+from .conv import SparseConv3d, SparseInverseConv3d
+from . import config
diff --git a/pixal3d/modules/sparse/conv/config.py b/pixal3d/modules/sparse/conv/config.py
new file mode 100644
index 0000000..ac08489
--- /dev/null
+++ b/pixal3d/modules/sparse/conv/config.py
@@ -0,0 +1,3 @@
+SPCONV_ALGO = 'auto'                                # 'auto', 'implicit_gemm', 'native'
+FLEX_GEMM_ALGO = 'masked_implicit_gemm_splitk'      # 'explicit_gemm', 'implicit_gemm', 'implicit_gemm_splitk', 'masked_implicit_gemm', 'masked_implicit_gemm_splitk'
+FLEX_GEMM_HASHMAP_RATIO = 2.0                       # Ratio of hashmap size to input size
diff --git a/pixal3d/modules/sparse/conv/conv.py b/pixal3d/modules/sparse/conv/conv.py
new file mode 100644
index 0000000..4c7d407
--- /dev/null
+++ b/pixal3d/modules/sparse/conv/conv.py
@@ -0,0 +1,30 @@
+from .. import config
+import importlib
+import torch
+import torch.nn as nn
+from .. import SparseTensor
+
+
+_backends = {}
+
+
+class SparseConv3d(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, padding=None, bias=True, indice_key=None):
+        super(SparseConv3d, self).__init__()
+        if config.CONV not in _backends:
+            _backends[config.CONV] = importlib.import_module(f'..conv_{config.CONV}', __name__)
+        _backends[config.CONV].sparse_conv3d_init(self, in_channels, out_channels, kernel_size, stride, dilation, padding, bias, indice_key)
+
+    def forward(self, x: SparseTensor) -> SparseTensor:
+        return _backends[config.CONV].sparse_conv3d_forward(self, x)
+
+
+class SparseInverseConv3d(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, bias=True, indice_key=None):
+        super(SparseInverseConv3d, self).__init__()
+        if config.CONV not in _backends:
+            _backends[config.CONV] = importlib.import_module(f'..conv_{config.CONV}', __name__)
+        _backends[config.CONV].sparse_inverse_conv3d_init(self, in_channels, out_channels, kernel_size, stride, dilation, bias, indice_key)
+
+    def forward(self, x: SparseTensor) -> SparseTensor:
+        return _backends[config.CONV].sparse_inverse_conv3d_forward(self, x)
diff --git a/pixal3d/modules/sparse/conv/conv_flex_gemm.py b/pixal3d/modules/sparse/conv/conv_flex_gemm.py
new file mode 100644
index 0000000..d256194
--- /dev/null
+++ b/pixal3d/modules/sparse/conv/conv_flex_gemm.py
@@ -0,0 +1,68 @@
+import math
+import torch
+import torch.nn as nn
+from .. import SparseTensor
+from . import config
+import flex_gemm
+from flex_gemm.ops.spconv import sparse_submanifold_conv3d
+
+
+def sparse_conv3d_init(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, padding=None, bias=True, indice_key=None):
+    assert stride == 1 and (padding is None), 'Currently flex_gemm implementation only support submanifold sparse convolution (stride=1, padding=None)'
+    
+    self.in_channels = in_channels
+    self.out_channels = out_channels
+    self.kernel_size = tuple(kernel_size) if isinstance(kernel_size, (list, tuple)) else (kernel_size, ) * 3
+    self.stride = tuple(stride) if isinstance(stride, (list, tuple)) else (stride, ) * 3
+    self.dilation = tuple(dilation) if isinstance(dilation, (list, tuple)) else (dilation, ) * 3
+
+    self.weight = nn.Parameter(torch.empty((out_channels, in_channels, *self.kernel_size)))
+    if bias:
+        self.bias = nn.Parameter(torch.empty(out_channels))
+    else:
+        self.register_parameter("bias", None)
+
+    # initialize parameters
+    torch.nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+    if self.bias is not None:
+        fan_in, _ = torch.nn.init._calculate_fan_in_and_fan_out(self.weight)
+        if fan_in != 0:
+            bound = 1 / math.sqrt(fan_in)
+            torch.nn.init.uniform_(self.bias, -bound, bound)
+
+    # Permute weight (Co, Ci, Kd, Kh, Kw) -> (Co, Kd, Kh, Kw, Ci)
+    self.weight = nn.Parameter(self.weight.permute(0, 2, 3, 4, 1).contiguous())
+
+
+def sparse_conv3d_forward(self, x: SparseTensor) -> SparseTensor:
+    flex_gemm.ops.spconv.set_algorithm(config.FLEX_GEMM_ALGO)
+    flex_gemm.ops.spconv.set_hashmap_ratio(config.FLEX_GEMM_HASHMAP_RATIO)
+
+    # check if neighbor map is already computed
+    Co, Kd, Kh, Kw, Ci = self.weight.shape
+    neighbor_cache_key = f'SubMConv3d_neighbor_cache_{Kw}x{Kh}x{Kd}_dilation{self.dilation}'
+    neighbor_cache = x.get_spatial_cache(neighbor_cache_key)
+    
+    out, neighbor_cache_ = sparse_submanifold_conv3d(
+        x.feats,
+        x.coords,
+        torch.Size([*x.shape, *x.spatial_shape]),
+        self.weight,
+        self.bias,
+        neighbor_cache,
+        self.dilation
+    )
+    
+    if neighbor_cache is None:
+        x.register_spatial_cache(neighbor_cache_key, neighbor_cache_)
+    
+    out = x.replace(out)
+    return out
+
+
+def sparse_inverse_conv3d_init(self, *args, **kwargs):
+    raise NotImplementedError('SparseInverseConv3d with flex_gemm is not implemented yet')
+
+
+def sparse_inverse_conv3d_forward(self, x: SparseTensor) -> SparseTensor:
+    raise NotImplementedError('SparseInverseConv3d with flex_gemm is not implemented yet')
diff --git a/pixal3d/modules/sparse/conv/conv_spconv.py b/pixal3d/modules/sparse/conv/conv_spconv.py
new file mode 100644
index 0000000..f709708
--- /dev/null
+++ b/pixal3d/modules/sparse/conv/conv_spconv.py
@@ -0,0 +1,73 @@
+import torch
+import torch.nn as nn
+from .. import SparseTensor
+from . import config
+import spconv.pytorch as spconv
+
+
+def sparse_conv3d_init(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, padding=None, bias=True, indice_key=None):
+    algo = None
+    if config.SPCONV_ALGO == 'native':
+        algo = spconv.ConvAlgo.Native
+    elif config.SPCONV_ALGO == 'implicit_gemm':
+        algo = spconv.ConvAlgo.MaskImplicitGemm
+    if stride == 1 and (padding is None):
+        self.conv = spconv.SubMConv3d(in_channels, out_channels, kernel_size, dilation=dilation, bias=bias, indice_key=indice_key, algo=algo)
+    else:
+        self.conv = spconv.SparseConv3d(in_channels, out_channels, kernel_size, stride=stride, dilation=dilation, padding=padding, bias=bias, indice_key=indice_key, algo=algo)
+    self.stride = tuple(stride) if isinstance(stride, (list, tuple)) else (stride, stride, stride)
+    self.padding = padding  
+
+
+def sparse_conv3d_forward(self, x: SparseTensor) -> SparseTensor:
+    spatial_changed = any(s != 1 for s in self.stride) or (self.padding is not None)
+    new_data = self.conv(x.data)
+    new_shape = [x.shape[0], self.conv.out_channels]
+    new_layout = None if spatial_changed else x.layout
+
+    if spatial_changed and (x.shape[0] != 1):
+        # spconv was non-1 stride will break the contiguous of the output tensor, sort by the coords
+        fwd = new_data.indices[:, 0].argsort()
+        bwd = torch.zeros_like(fwd).scatter_(0, fwd, torch.arange(fwd.shape[0], device=fwd.device))
+        sorted_feats = new_data.features[fwd]
+        sorted_coords = new_data.indices[fwd]
+        unsorted_data = new_data
+        new_data = spconv.SparseConvTensor(sorted_feats, sorted_coords, unsorted_data.spatial_shape, unsorted_data.batch_size)  # type: ignore
+
+    out = SparseTensor(
+        new_data, shape=torch.Size(new_shape), layout=new_layout,
+        scale=tuple([s * stride for s, stride in zip(x._scale, self.stride)]),
+        spatial_cache=x._spatial_cache,
+    )
+
+    if spatial_changed and (x.shape[0] != 1):
+        out.register_spatial_cache(f'conv_{self.stride}_unsorted_data', unsorted_data)
+        out.register_spatial_cache(f'conv_{self.stride}_sort_bwd', bwd)
+
+    return out
+
+
+def sparse_inverse_conv3d_init(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, bias=True, indice_key=None):
+    self.conv = spconv.SparseInverseConv3d(in_channels, out_channels, kernel_size, bias=bias, indice_key=indice_key)
+    self.stride = tuple(stride) if isinstance(stride, (list, tuple)) else (stride, stride, stride)
+
+
+def sparse_inverse_conv3d_forward(self, x: SparseTensor) -> SparseTensor:
+    spatial_changed = any(s != 1 for s in self.stride)
+    if spatial_changed:
+        # recover the original spconv order
+        data = x.get_spatial_cache(f'conv_{self.stride}_unsorted_data')
+        bwd = x.get_spatial_cache(f'conv_{self.stride}_sort_bwd')
+        data = data.replace_feature(x.feats[bwd])
+    else:
+        data = x.data
+
+    new_data = self.conv(data)
+    new_shape = [x.shape[0], self.conv.out_channels]
+    new_layout = None if spatial_changed else x.layout
+    out = SparseTensor(
+        new_data, shape=torch.Size(new_shape), layout=new_layout,
+        scale=tuple([s // stride for s, stride in zip(x._scale, self.stride)]),
+        spatial_cache=x._spatial_cache,
+    )
+    return out
diff --git a/pixal3d/modules/sparse/conv/conv_torchsparse.py b/pixal3d/modules/sparse/conv/conv_torchsparse.py
new file mode 100644
index 0000000..5234bd1
--- /dev/null
+++ b/pixal3d/modules/sparse/conv/conv_torchsparse.py
@@ -0,0 +1,30 @@
+import torch
+import torch.nn as nn
+from .. import SparseTensor
+import torchsparse
+
+
+def sparse_conv3d_init(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, padding=None, bias=True, indice_key=None):
+    self.conv = torchsparse.nn.Conv3d(in_channels, out_channels, kernel_size, stride, 0, dilation, bias)
+
+
+def sparse_conv3d_forward(self, x: SparseTensor) -> SparseTensor:
+    out = self.conv(x.data)
+    new_shape = [x.shape[0], self.conv.out_channels]
+    out = SparseTensor(out, shape=torch.Size(new_shape), layout=x.layout if all(s == 1 for s in self.conv.stride) else None)
+    out._spatial_cache = x._spatial_cache
+    out._scale = tuple([s * stride for s, stride in zip(x._scale, self.conv.stride)])
+    return out
+
+
+def sparse_inverse_conv3d_init(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, bias=True, indice_key=None):
+    self.conv = torchsparse.nn.Conv3d(in_channels, out_channels, kernel_size, stride, 0, dilation, bias, transposed=True)
+
+
+def sparse_inverse_conv3d_forward(self, x: SparseTensor) -> SparseTensor:
+    out = self.conv(x.data)        
+    new_shape = [x.shape[0], self.conv.out_channels]
+    out = SparseTensor(out, shape=torch.Size(new_shape), layout=x.layout if all(s == 1 for s in self.conv.stride) else None)
+    out._spatial_cache = x._spatial_cache
+    out._scale = tuple([s / stride for s, stride in zip(x._scale, self.conv.stride)])
+    return out
diff --git a/pixal3d/modules/sparse/linear.py b/pixal3d/modules/sparse/linear.py
new file mode 100644
index 0000000..4431770
--- /dev/null
+++ b/pixal3d/modules/sparse/linear.py
@@ -0,0 +1,15 @@
+import torch
+import torch.nn as nn
+from . import VarLenTensor
+
+__all__ = [
+    'SparseLinear'
+]
+
+
+class SparseLinear(nn.Linear):
+    def __init__(self, in_features, out_features, bias=True):
+        super(SparseLinear, self).__init__(in_features, out_features, bias)
+
+    def forward(self, input: VarLenTensor) -> VarLenTensor:
+        return input.replace(super().forward(input.feats))
diff --git a/pixal3d/modules/sparse/nonlinearity.py b/pixal3d/modules/sparse/nonlinearity.py
new file mode 100644
index 0000000..950e5c0
--- /dev/null
+++ b/pixal3d/modules/sparse/nonlinearity.py
@@ -0,0 +1,35 @@
+import torch
+import torch.nn as nn
+from . import VarLenTensor
+
+__all__ = [
+    'SparseReLU',
+    'SparseSiLU',
+    'SparseGELU',
+    'SparseActivation'
+]
+
+
+class SparseReLU(nn.ReLU):
+    def forward(self, input: VarLenTensor) -> VarLenTensor:
+        return input.replace(super().forward(input.feats))
+    
+
+class SparseSiLU(nn.SiLU):
+    def forward(self, input: VarLenTensor) -> VarLenTensor:
+        return input.replace(super().forward(input.feats))
+
+
+class SparseGELU(nn.GELU):
+    def forward(self, input: VarLenTensor) -> VarLenTensor:
+        return input.replace(super().forward(input.feats))
+
+
+class SparseActivation(nn.Module):
+    def __init__(self, activation: nn.Module):
+        super().__init__()
+        self.activation = activation
+
+    def forward(self, input: VarLenTensor) -> VarLenTensor:
+        return input.replace(self.activation(input.feats))
+    
diff --git a/pixal3d/modules/sparse/norm.py b/pixal3d/modules/sparse/norm.py
new file mode 100644
index 0000000..9571120
--- /dev/null
+++ b/pixal3d/modules/sparse/norm.py
@@ -0,0 +1,64 @@
+import torch
+import torch.nn as nn
+from ..utils import manual_cast
+from . import VarLenTensor
+from . import config
+
+__all__ = [
+    'SparseGroupNorm',
+    'SparseLayerNorm',
+    'SparseGroupNorm32',
+    'SparseLayerNorm32',
+]
+
+
+class SparseGroupNorm(nn.GroupNorm):
+    def __init__(self, num_groups, num_channels, eps=1e-5, affine=True):
+        super(SparseGroupNorm, self).__init__(num_groups, num_channels, eps, affine)
+
+    def forward(self, input: VarLenTensor) -> VarLenTensor:
+        nfeats = torch.zeros_like(input.feats)
+        for k in range(input.shape[0]):
+            bfeats = input.feats[input.layout[k]]
+            bfeats = bfeats.permute(1, 0).reshape(1, input.shape[1], -1)
+            bfeats = super().forward(bfeats)
+            bfeats = bfeats.reshape(input.shape[1], -1).permute(1, 0)
+            nfeats[input.layout[k]] = bfeats
+        return input.replace(nfeats)
+
+
+class SparseLayerNorm(nn.LayerNorm):
+    def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True):
+        super(SparseLayerNorm, self).__init__(normalized_shape, eps, elementwise_affine)
+
+    def forward(self, input: VarLenTensor) -> VarLenTensor:
+        nfeats = torch.zeros_like(input.feats)
+        for k in range(input.shape[0]):
+            bfeats = input.feats[input.layout[k]]
+            bfeats = bfeats.permute(1, 0).reshape(1, input.shape[1], -1)
+            bfeats = super().forward(bfeats)
+            bfeats = bfeats.reshape(input.shape[1], -1).permute(1, 0)
+            nfeats[input.layout[k]] = bfeats
+        return input.replace(nfeats)
+
+
+class SparseGroupNorm32(SparseGroupNorm):
+    """
+    A GroupNorm layer that converts to float32 before the forward pass.
+    """
+    def forward(self, x: VarLenTensor) -> VarLenTensor:
+        x_dtype = x.dtype
+        x = manual_cast(x, torch.float32)
+        o = super().forward(x)
+        return manual_cast(o, x_dtype)
+
+
+class SparseLayerNorm32(SparseLayerNorm):
+    """
+    A LayerNorm layer that converts to float32 before the forward pass.
+    """
+    def forward(self, x: VarLenTensor) -> VarLenTensor:
+        x_dtype = x.dtype
+        x = manual_cast(x, torch.float32)
+        o = super().forward(x)
+        return manual_cast(o, x_dtype)
diff --git a/pixal3d/modules/sparse/spatial/__init__.py b/pixal3d/modules/sparse/spatial/__init__.py
new file mode 100644
index 0000000..e27425f
--- /dev/null
+++ b/pixal3d/modules/sparse/spatial/__init__.py
@@ -0,0 +1,2 @@
+from .basic import *
+from .spatial2channel import *
diff --git a/pixal3d/modules/sparse/spatial/basic.py b/pixal3d/modules/sparse/spatial/basic.py
new file mode 100644
index 0000000..eaeb8af
--- /dev/null
+++ b/pixal3d/modules/sparse/spatial/basic.py
@@ -0,0 +1,109 @@
+from typing import *
+import torch
+import torch.nn as nn
+from .. import SparseTensor
+
+__all__ = [
+    'SparseDownsample',
+    'SparseUpsample',
+]
+
+
+class SparseDownsample(nn.Module):
+    """
+    Downsample a sparse tensor by a factor of `factor`.
+    Implemented as average pooling.
+    """
+    def __init__(self, factor: int, mode: Literal['mean', 'max'] = 'mean'):
+        super(SparseDownsample, self).__init__()
+        self.factor = factor
+        self.mode = mode
+        assert self.mode in ['mean', 'max'], f'Invalid mode: {self.mode}'
+
+    def forward(self, x: SparseTensor) -> SparseTensor:
+        cache = x.get_spatial_cache(f'downsample_{self.factor}')
+        if cache is None:
+            DIM = x.coords.shape[-1] - 1
+
+            coord = list(x.coords.unbind(dim=-1))
+            for i in range(DIM):
+                coord[i+1] = coord[i+1] // self.factor
+
+            MAX = [(s + self.factor - 1) // self.factor for s in x.spatial_shape]
+            OFFSET = torch.cumprod(torch.tensor(MAX[::-1]), 0).tolist()[::-1] + [1]
+            code = sum([c * o for c, o in zip(coord, OFFSET)])
+            code, idx = code.unique(return_inverse=True)
+
+            new_coords = torch.stack(
+                [code // OFFSET[0]] +
+                [(code // OFFSET[i+1]) % MAX[i] for i in range(DIM)],
+                dim=-1
+            )
+        else:
+            new_coords, idx = cache
+            
+        new_feats = torch.scatter_reduce(
+            torch.zeros(new_coords.shape[0], x.feats.shape[1], device=x.feats.device, dtype=x.feats.dtype),
+            dim=0,
+            index=idx.unsqueeze(1).expand(-1, x.feats.shape[1]),
+            src=x.feats,
+            reduce=self.mode,
+            include_self=False,
+        )
+        out = SparseTensor(new_feats, new_coords, x._shape)
+        out._scale = tuple([s * self.factor for s in x._scale])
+        out._spatial_cache = x._spatial_cache
+        
+        if cache is None:
+            x.register_spatial_cache(f'downsample_{self.factor}', (new_coords, idx))
+            out.register_spatial_cache(f'upsample_{self.factor}', (x.coords, idx))
+            out.register_spatial_cache(f'shape', torch.Size(MAX))
+            if self.training:
+                subidx = x.coords[:, 1:] % self.factor
+                subidx = sum([subidx[..., i] * self.factor ** i for i in range(DIM)])
+                subdivision = torch.zeros((new_coords.shape[0], self.factor ** DIM), device=x.device, dtype=torch.bool)
+                subdivision[idx, subidx] = True
+                out.register_spatial_cache(f'subdivision', subdivision)
+
+        return out
+
+
+class SparseUpsample(nn.Module):
+    """
+    Upsample a sparse tensor by a factor of `factor`.
+    Implemented as nearest neighbor interpolation.
+    """
+    def __init__(
+        self, factor: int
+    ):
+        super(SparseUpsample, self).__init__()
+        self.factor = factor
+
+    def forward(self, x: SparseTensor, subdivision: Optional[SparseTensor] = None) -> SparseTensor:
+        DIM = x.coords.shape[-1] - 1
+
+        cache = x.get_spatial_cache(f'upsample_{self.factor}')
+        if cache is None:
+            if subdivision is None:
+                raise ValueError('Cache not found. Provide subdivision tensor or pair SparseUpsample with SparseDownsample.')
+            else:
+                sub = subdivision.feats
+                N_leaf = sub.sum(dim=-1)
+                subidx = sub.nonzero()[:, -1]
+                new_coords = x.coords.clone().detach()
+                new_coords[:, 1:] *= self.factor
+                new_coords = torch.repeat_interleave(new_coords, N_leaf, dim=0, output_size=subidx.shape[0])
+                for i in range(DIM):
+                    new_coords[:, i+1] += subidx // self.factor ** i % self.factor
+                idx = torch.repeat_interleave(torch.arange(x.coords.shape[0], device=x.device), N_leaf, dim=0, output_size=subidx.shape[0])
+        else:
+            new_coords, idx = cache
+            
+        new_feats = x.feats[idx]
+        out = SparseTensor(new_feats, new_coords, x._shape)
+        out._scale = tuple([s / self.factor for s in x._scale])
+        if cache is not None:           # only keep cache when subdiv following it
+            out._spatial_cache = x._spatial_cache
+        
+        return out
+ 
\ No newline at end of file
diff --git a/pixal3d/modules/sparse/spatial/spatial2channel.py b/pixal3d/modules/sparse/spatial/spatial2channel.py
new file mode 100644
index 0000000..577f36d
--- /dev/null
+++ b/pixal3d/modules/sparse/spatial/spatial2channel.py
@@ -0,0 +1,93 @@
+from typing import *
+import torch
+import torch.nn as nn
+from .. import SparseTensor
+
+
+class SparseSpatial2Channel(nn.Module):
+    """
+    Downsample a sparse tensor by a factor of `factor`.
+    Implemented as rearranging its features from spatial to channel.
+    """
+    def __init__(self, factor: int = 2):
+        super(SparseSpatial2Channel, self).__init__()
+        self.factor = factor
+
+    def forward(self, x: SparseTensor) -> SparseTensor:
+        DIM = x.coords.shape[-1] - 1
+        cache = x.get_spatial_cache(f'spatial2channel_{self.factor}')
+        if cache is None:
+            coord = list(x.coords.unbind(dim=-1))
+            for i in range(DIM):
+                coord[i+1] = coord[i+1] // self.factor
+            subidx = x.coords[:, 1:] % self.factor
+            subidx = sum([subidx[..., i] * self.factor ** i for i in range(DIM)])
+
+            MAX = [(s + self.factor - 1) // self.factor for s in x.spatial_shape]
+            OFFSET = torch.cumprod(torch.tensor(MAX[::-1]), 0).tolist()[::-1] + [1]
+            code = sum([c * o for c, o in zip(coord, OFFSET)])
+            code, idx = code.unique(return_inverse=True)
+
+            new_coords = torch.stack(
+                [code // OFFSET[0]] +
+                [(code // OFFSET[i+1]) % MAX[i] for i in range(DIM)],
+                dim=-1
+            )
+        else:
+            new_coords, idx, subidx = cache
+            
+        new_feats = torch.zeros(new_coords.shape[0] * self.factor ** DIM, x.feats.shape[1], device=x.feats.device, dtype=x.feats.dtype)
+        new_feats[idx * self.factor ** DIM + subidx] = x.feats
+
+        out = SparseTensor(new_feats.reshape(new_coords.shape[0], -1), new_coords, None if x._shape is None else torch.Size([x._shape[0], x._shape[1] * self.factor ** DIM]))
+        out._scale = tuple([s * self.factor for s in x._scale])
+        out._spatial_cache = x._spatial_cache
+        
+        if cache is None:
+            x.register_spatial_cache(f'spatial2channel_{self.factor}', (new_coords, idx, subidx))
+            out.register_spatial_cache(f'channel2spatial_{self.factor}', (x.coords, idx, subidx))
+            out.register_spatial_cache(f'shape', torch.Size(MAX))
+            if self.training:
+                subdivision = torch.zeros((new_coords.shape[0], self.factor ** DIM), device=x.device, dtype=torch.bool)
+                subdivision[idx, subidx] = True
+                out.register_spatial_cache(f'subdivision', subdivision)
+                
+        return out
+
+
+class SparseChannel2Spatial(nn.Module):
+    """
+    Upsample a sparse tensor by a factor of `factor`.
+    Implemented as rearranging its features from channel to spatial.
+    """
+    def __init__(self, factor: int = 2):
+        super(SparseChannel2Spatial, self).__init__()
+        self.factor = factor
+
+    def forward(self, x: SparseTensor, subdivision: Optional[SparseTensor] = None) -> SparseTensor:
+        DIM = x.coords.shape[-1] - 1
+
+        cache = x.get_spatial_cache(f'channel2spatial_{self.factor}')
+        if cache is None:
+            if subdivision is None:
+                raise ValueError('Cache not found. Provide subdivision tensor or pair SparseChannel2Spatial with SparseSpatial2Channel.')
+            else:
+                sub = subdivision.feats         # [N, self.factor ** DIM]
+                N_leaf = sub.sum(dim=-1)        # [N]
+                subidx = sub.nonzero()[:, -1]
+                new_coords = x.coords.clone().detach()
+                new_coords[:, 1:] *= self.factor
+                new_coords = torch.repeat_interleave(new_coords, N_leaf, dim=0, output_size=subidx.shape[0])
+                for i in range(DIM):
+                    new_coords[:, i+1] += subidx // self.factor ** i % self.factor
+                idx = torch.repeat_interleave(torch.arange(x.coords.shape[0], device=x.device), N_leaf, dim=0, output_size=subidx.shape[0])
+        else:
+            new_coords, idx, subidx = cache
+
+        x_feats = x.feats.reshape(x.feats.shape[0] * self.factor ** DIM, -1)
+        new_feats = x_feats[idx * self.factor ** DIM + subidx]
+        out = SparseTensor(new_feats, new_coords, None if x._shape is None else torch.Size([x._shape[0], x._shape[1] // self.factor ** DIM]))
+        out._scale = tuple([s / self.factor for s in x._scale])
+        if cache is not None:           # only keep cache when subdiv following it
+            out._spatial_cache = x._spatial_cache
+        return out
diff --git a/pixal3d/modules/sparse/transformer/__init__.py b/pixal3d/modules/sparse/transformer/__init__.py
new file mode 100644
index 0000000..b08b0d4
--- /dev/null
+++ b/pixal3d/modules/sparse/transformer/__init__.py
@@ -0,0 +1,2 @@
+from .blocks import *
+from .modulated import *
\ No newline at end of file
diff --git a/pixal3d/modules/sparse/transformer/blocks.py b/pixal3d/modules/sparse/transformer/blocks.py
new file mode 100644
index 0000000..9d1ec60
--- /dev/null
+++ b/pixal3d/modules/sparse/transformer/blocks.py
@@ -0,0 +1,145 @@
+from typing import *
+import torch
+import torch.nn as nn
+from ..basic import VarLenTensor, SparseTensor
+from ..linear import SparseLinear
+from ..nonlinearity import SparseGELU
+from ..attention import SparseMultiHeadAttention
+from ...norm import LayerNorm32
+
+
+class SparseFeedForwardNet(nn.Module):
+    def __init__(self, channels: int, mlp_ratio: float = 4.0):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            SparseLinear(channels, int(channels * mlp_ratio)),
+            SparseGELU(approximate="tanh"),
+            SparseLinear(int(channels * mlp_ratio), channels),
+        )
+
+    def forward(self, x: VarLenTensor) -> VarLenTensor:
+        return self.mlp(x)
+
+
+class SparseTransformerBlock(nn.Module):
+    """
+    Sparse Transformer block (MSA + FFN).
+    """
+    def __init__(
+        self,
+        channels: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        attn_mode: Literal["full", "swin"] = "full",
+        window_size: Optional[int] = None,
+        shift_window: Optional[Tuple[int, int, int]] = None,
+        use_checkpoint: bool = False,
+        use_rope: bool = False,
+        rope_freq: Tuple[int, int] = (1.0, 10000.0), 
+        qk_rms_norm: bool = False,
+        qkv_bias: bool = True,
+        ln_affine: bool = False,
+    ):
+        super().__init__()
+        self.use_checkpoint = use_checkpoint
+        self.norm1 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6)
+        self.norm2 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6)
+        self.attn = SparseMultiHeadAttention(
+            channels,
+            num_heads=num_heads,
+            attn_mode=attn_mode,
+            window_size=window_size,
+            shift_window=shift_window,
+            qkv_bias=qkv_bias,
+            use_rope=use_rope,
+            rope_freq=rope_freq,
+            qk_rms_norm=qk_rms_norm,
+        )
+        self.mlp = SparseFeedForwardNet(
+            channels,
+            mlp_ratio=mlp_ratio,
+        )
+
+    def _forward(self, x: SparseTensor) -> SparseTensor:
+        h = x.replace(self.norm1(x.feats))
+        h = self.attn(h)
+        x = x + h
+        h = x.replace(self.norm2(x.feats))
+        h = self.mlp(h)
+        x = x + h
+        return x
+
+    def forward(self, x: SparseTensor) -> SparseTensor:
+        if self.use_checkpoint:
+            return torch.utils.checkpoint.checkpoint(self._forward, x, use_reentrant=False)
+        else:
+            return self._forward(x)
+
+
+class SparseTransformerCrossBlock(nn.Module):
+    """
+    Sparse Transformer cross-attention block (MSA + MCA + FFN).
+    """
+    def __init__(
+        self,
+        channels: int,
+        ctx_channels: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        attn_mode: Literal["full", "swin"] = "full",
+        window_size: Optional[int] = None,
+        shift_window: Optional[Tuple[int, int, int]] = None,
+        use_checkpoint: bool = False,
+        use_rope: bool = False,
+        qk_rms_norm: bool = False,
+        qk_rms_norm_cross: bool = False,
+        qkv_bias: bool = True,
+        ln_affine: bool = False,
+    ):
+        super().__init__()
+        self.use_checkpoint = use_checkpoint
+        self.norm1 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6)
+        self.norm2 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6)
+        self.norm3 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6)
+        self.self_attn = SparseMultiHeadAttention(
+            channels,
+            num_heads=num_heads,
+            type="self",
+            attn_mode=attn_mode,
+            window_size=window_size,
+            shift_window=shift_window,
+            qkv_bias=qkv_bias,
+            use_rope=use_rope,
+            qk_rms_norm=qk_rms_norm,
+        )
+        self.cross_attn = SparseMultiHeadAttention(
+            channels,
+            ctx_channels=ctx_channels,
+            num_heads=num_heads,
+            type="cross",
+            attn_mode="full",
+            qkv_bias=qkv_bias,
+            qk_rms_norm=qk_rms_norm_cross,
+        )
+        self.mlp = SparseFeedForwardNet(
+            channels,
+            mlp_ratio=mlp_ratio,
+        )
+
+    def _forward(self, x: SparseTensor, context: Union[torch.Tensor, VarLenTensor]) -> SparseTensor:
+        h = x.replace(self.norm1(x.feats))
+        h = self.self_attn(h)
+        x = x + h
+        h = x.replace(self.norm2(x.feats))
+        h = self.cross_attn(h, context)
+        x = x + h
+        h = x.replace(self.norm3(x.feats))
+        h = self.mlp(h)
+        x = x + h
+        return x
+
+    def forward(self, x: SparseTensor, context: Union[torch.Tensor, VarLenTensor]) -> SparseTensor:
+        if self.use_checkpoint:
+            return torch.utils.checkpoint.checkpoint(self._forward, x, context, use_reentrant=False)
+        else:
+            return self._forward(x, context)
diff --git a/pixal3d/modules/sparse/transformer/modulated.py b/pixal3d/modules/sparse/transformer/modulated.py
new file mode 100644
index 0000000..fc8fa56
--- /dev/null
+++ b/pixal3d/modules/sparse/transformer/modulated.py
@@ -0,0 +1,205 @@
+from typing import *
+import torch
+import torch.nn as nn
+from ..basic import VarLenTensor, SparseTensor
+from ..attention import SparseMultiHeadAttention, SparseProjectAttention, SparseGatedProjectAttention
+from ...norm import LayerNorm32
+from .blocks import SparseFeedForwardNet
+
+
+class ModulatedSparseTransformerBlock(nn.Module):
+    """
+    Sparse Transformer block (MSA + FFN) with adaptive layer norm conditioning.
+    """
+    def __init__(
+        self,
+        channels: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        attn_mode: Literal["full", "swin"] = "full",
+        window_size: Optional[int] = None,
+        shift_window: Optional[Tuple[int, int, int]] = None,
+        use_checkpoint: bool = False,
+        use_rope: bool = False,
+        rope_freq: Tuple[float, float] = (1.0, 10000.0),
+        qk_rms_norm: bool = False,
+        qkv_bias: bool = True,
+        share_mod: bool = False,
+    ):
+        super().__init__()
+        self.use_checkpoint = use_checkpoint
+        self.share_mod = share_mod
+        self.norm1 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)
+        self.norm2 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)
+        self.attn = SparseMultiHeadAttention(
+            channels,
+            num_heads=num_heads,
+            attn_mode=attn_mode,
+            window_size=window_size,
+            shift_window=shift_window,
+            qkv_bias=qkv_bias,
+            use_rope=use_rope,
+            rope_freq=rope_freq,
+            qk_rms_norm=qk_rms_norm,
+        )
+        self.mlp = SparseFeedForwardNet(
+            channels,
+            mlp_ratio=mlp_ratio,
+        )
+        if not share_mod:
+            self.adaLN_modulation = nn.Sequential(
+                nn.SiLU(),
+                nn.Linear(channels, 6 * channels, bias=True)
+            )
+        else:
+            self.modulation = nn.Parameter(torch.randn(6 * channels) / channels ** 0.5)
+
+    def _forward(self, x: SparseTensor, mod: torch.Tensor) -> SparseTensor:
+        if self.share_mod:
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (self.modulation + mod).type(mod.dtype).chunk(6, dim=1)
+        else:
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(mod).chunk(6, dim=1)
+        h = x.replace(self.norm1(x.feats))
+        h = h * (1 + scale_msa) + shift_msa
+        h = self.attn(h)
+        h = h * gate_msa
+        x = x + h
+        h = x.replace(self.norm2(x.feats))
+        h = h * (1 + scale_mlp) + shift_mlp
+        h = self.mlp(h)
+        h = h * gate_mlp
+        x = x + h
+        return x
+
+    def forward(self, x: SparseTensor, mod: torch.Tensor) -> SparseTensor:
+        if self.use_checkpoint:
+            return torch.utils.checkpoint.checkpoint(self._forward, x, mod, use_reentrant=False)
+        else:
+            return self._forward(x, mod)
+
+
+class ModulatedSparseTransformerCrossBlock(nn.Module):
+    """
+    Sparse Transformer cross-attention block (MSA + MCA + FFN) with adaptive layer norm conditioning.
+    
+    Supports two image attention modes:
+    - "cross": Standard cross-attention with image features
+    - "proj": Projection-based attention with view-aligned features
+    """
+    def __init__(
+        self,
+        channels: int,
+        ctx_channels: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        attn_mode: Literal["full", "swin"] = "full",
+        window_size: Optional[int] = None,
+        shift_window: Optional[Tuple[int, int, int]] = None,
+        use_checkpoint: bool = False,
+        use_rope: bool = False,
+        rope_freq: Tuple[float, float] = (1.0, 10000.0),
+        qk_rms_norm: bool = False,
+        qk_rms_norm_cross: bool = False,
+        qkv_bias: bool = True,
+        share_mod: bool = False,
+        image_attn_mode: Literal["cross", "proj", "gated_proj"] = "cross",
+        proj_in_channels: Optional[int] = None,
+        vae_in_channels: Optional[int] = None,
+    ):
+        super().__init__()
+        self.use_checkpoint = use_checkpoint
+        self.share_mod = share_mod
+        self.image_attn_mode = image_attn_mode
+        
+        self.norm1 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)
+        self.norm2 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
+        self.norm3 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)
+        self.self_attn = SparseMultiHeadAttention(
+            channels,
+            num_heads=num_heads,
+            type="self",
+            attn_mode=attn_mode,
+            window_size=window_size,
+            shift_window=shift_window,
+            qkv_bias=qkv_bias,
+            use_rope=use_rope,
+            rope_freq=rope_freq,
+            qk_rms_norm=qk_rms_norm,
+        )
+        
+        # Build cross attention based on mode
+        if image_attn_mode == "cross":
+            self.cross_attn = SparseMultiHeadAttention(
+                channels,
+                ctx_channels=ctx_channels,
+                num_heads=num_heads,
+                type="cross",
+                attn_mode="full",
+                qkv_bias=qkv_bias,
+                qk_rms_norm=qk_rms_norm_cross,
+            )
+        elif image_attn_mode == "proj":
+            _proj_in = proj_in_channels if proj_in_channels is not None else ctx_channels
+            cross_attn_block = SparseMultiHeadAttention(
+                channels,
+                ctx_channels=ctx_channels,
+                num_heads=num_heads,
+                type="cross",
+                attn_mode="full",
+                qkv_bias=qkv_bias,
+                qk_rms_norm=qk_rms_norm_cross,
+            )
+            self.cross_attn = SparseProjectAttention(cross_attn_block, channels, _proj_in)
+        elif image_attn_mode == "gated_proj":
+            _dino_in = proj_in_channels if proj_in_channels is not None else ctx_channels
+            _vae_in = vae_in_channels if vae_in_channels is not None else 16
+            cross_attn_block = SparseMultiHeadAttention(
+                channels,
+                ctx_channels=ctx_channels,
+                num_heads=num_heads,
+                type="cross",
+                attn_mode="full",
+                qkv_bias=qkv_bias,
+                qk_rms_norm=qk_rms_norm_cross,
+            )
+            self.cross_attn = SparseGatedProjectAttention(cross_attn_block, channels, _dino_in, _vae_in)
+        else:
+            raise ValueError(f"Unknown image attention mode: {image_attn_mode}")
+            
+        self.mlp = SparseFeedForwardNet(
+            channels,
+            mlp_ratio=mlp_ratio,
+        )
+        if not share_mod:
+            self.adaLN_modulation = nn.Sequential(
+                nn.SiLU(),
+                nn.Linear(channels, 6 * channels, bias=True)
+            )
+        else:
+            self.modulation = nn.Parameter(torch.randn(6 * channels) / channels ** 0.5)
+
+    def _forward(self, x: SparseTensor, mod: torch.Tensor, context: Union[torch.Tensor, VarLenTensor]) -> SparseTensor:
+        if self.share_mod:
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (self.modulation + mod).type(mod.dtype).chunk(6, dim=1)
+        else:
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(mod).chunk(6, dim=1)
+        h = x.replace(self.norm1(x.feats))
+        h = h * (1 + scale_msa) + shift_msa
+        h = self.self_attn(h)
+        h = h * gate_msa
+        x = x + h
+        h = x.replace(self.norm2(x.feats))
+        h = self.cross_attn(h, context)
+        x = x + h
+        h = x.replace(self.norm3(x.feats))
+        h = h * (1 + scale_mlp) + shift_mlp
+        h = self.mlp(h)
+        h = h * gate_mlp
+        x = x + h
+        return x
+
+    def forward(self, x: SparseTensor, mod: torch.Tensor, context: Union[torch.Tensor, VarLenTensor]) -> SparseTensor:
+        if self.use_checkpoint:
+            return torch.utils.checkpoint.checkpoint(self._forward, x, mod, context, use_reentrant=False)
+        else:
+            return self._forward(x, mod, context)
diff --git a/pixal3d/modules/spatial.py b/pixal3d/modules/spatial.py
new file mode 100644
index 0000000..79e268d
--- /dev/null
+++ b/pixal3d/modules/spatial.py
@@ -0,0 +1,48 @@
+import torch
+
+
+def pixel_shuffle_3d(x: torch.Tensor, scale_factor: int) -> torch.Tensor:
+    """
+    3D pixel shuffle.
+    """
+    B, C, H, W, D = x.shape
+    C_ = C // scale_factor**3
+    x = x.reshape(B, C_, scale_factor, scale_factor, scale_factor, H, W, D)
+    x = x.permute(0, 1, 5, 2, 6, 3, 7, 4)
+    x = x.reshape(B, C_, H*scale_factor, W*scale_factor, D*scale_factor)
+    return x
+
+
+def patchify(x: torch.Tensor, patch_size: int):
+    """
+    Patchify a tensor.
+
+    Args:
+        x (torch.Tensor): (N, C, *spatial) tensor
+        patch_size (int): Patch size
+    """
+    DIM = x.dim() - 2
+    for d in range(2, DIM + 2):
+        assert x.shape[d] % patch_size == 0, f"Dimension {d} of input tensor must be divisible by patch size, got {x.shape[d]} and {patch_size}"
+
+    x = x.reshape(*x.shape[:2], *sum([[x.shape[d] // patch_size, patch_size] for d in range(2, DIM + 2)], []))
+    x = x.permute(0, 1, *([2 * i + 3 for i in range(DIM)] + [2 * i + 2 for i in range(DIM)]))
+    x = x.reshape(x.shape[0], x.shape[1] * (patch_size ** DIM), *(x.shape[-DIM:]))
+    return x
+
+
+def unpatchify(x: torch.Tensor, patch_size: int):
+    """
+    Unpatchify a tensor.
+
+    Args:
+        x (torch.Tensor): (N, C, *spatial) tensor
+        patch_size (int): Patch size
+    """
+    DIM = x.dim() - 2
+    assert x.shape[1] % (patch_size ** DIM) == 0, f"Second dimension of input tensor must be divisible by patch size to unpatchify, got {x.shape[1]} and {patch_size ** DIM}"
+
+    x = x.reshape(x.shape[0], x.shape[1] // (patch_size ** DIM), *([patch_size] * DIM), *(x.shape[-DIM:]))
+    x = x.permute(0, 1, *(sum([[2 + DIM + i, 2 + i] for i in range(DIM)], [])))
+    x = x.reshape(x.shape[0], x.shape[1], *[x.shape[2 + 2 * i] * patch_size for i in range(DIM)])
+    return x
diff --git a/pixal3d/modules/transformer/__init__.py b/pixal3d/modules/transformer/__init__.py
new file mode 100644
index 0000000..b08b0d4
--- /dev/null
+++ b/pixal3d/modules/transformer/__init__.py
@@ -0,0 +1,2 @@
+from .blocks import *
+from .modulated import *
\ No newline at end of file
diff --git a/pixal3d/modules/transformer/blocks.py b/pixal3d/modules/transformer/blocks.py
new file mode 100644
index 0000000..fb6f5eb
--- /dev/null
+++ b/pixal3d/modules/transformer/blocks.py
@@ -0,0 +1,186 @@
+from typing import *
+import torch
+import torch.nn as nn
+from ..attention import MultiHeadAttention
+from ..norm import LayerNorm32
+
+
+class AbsolutePositionEmbedder(nn.Module):
+    """
+    Embeds spatial positions into vector representations.
+    """
+    def __init__(self, channels: int, in_channels: int = 3):
+        super().__init__()
+        self.channels = channels
+        self.in_channels = in_channels
+        self.freq_dim = channels // in_channels // 2
+        self.freqs = torch.arange(self.freq_dim, dtype=torch.float32) / self.freq_dim
+        self.freqs = 1.0 / (10000 ** self.freqs)
+        
+    def _sin_cos_embedding(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Create sinusoidal position embeddings.
+
+        Args:
+            x: a 1-D Tensor of N indices
+
+        Returns:
+            an (N, D) Tensor of positional embeddings.
+        """
+        self.freqs = self.freqs.to(x.device)
+        out = torch.outer(x, self.freqs)
+        out = torch.cat([torch.sin(out), torch.cos(out)], dim=-1)
+        return out
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x (torch.Tensor): (N, D) tensor of spatial positions
+        """
+        N, D = x.shape
+        assert D == self.in_channels, "Input dimension must match number of input channels"
+        embed = self._sin_cos_embedding(x.reshape(-1))
+        embed = embed.reshape(N, -1)
+        if embed.shape[1] < self.channels:
+            embed = torch.cat([embed, torch.zeros(N, self.channels - embed.shape[1], device=embed.device)], dim=-1)
+        return embed
+
+
+class FeedForwardNet(nn.Module):
+    def __init__(self, channels: int, mlp_ratio: float = 4.0):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(channels, int(channels * mlp_ratio)),
+            nn.GELU(approximate="tanh"),
+            nn.Linear(int(channels * mlp_ratio), channels),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.mlp(x)
+
+
+class TransformerBlock(nn.Module):
+    """
+    Transformer block (MSA + FFN).
+    """
+    def __init__(
+        self,
+        channels: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        attn_mode: Literal["full", "windowed"] = "full",
+        window_size: Optional[int] = None,
+        shift_window: Optional[int] = None,
+        use_checkpoint: bool = False,
+        use_rope: bool = False,
+        rope_freq: Tuple[int, int] = (1.0, 10000.0),
+        qk_rms_norm: bool = False,
+        qkv_bias: bool = True,
+        ln_affine: bool = True,
+    ):
+        super().__init__()
+        self.use_checkpoint = use_checkpoint
+        self.norm1 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6)
+        self.norm2 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6)
+        self.attn = MultiHeadAttention(
+            channels,
+            num_heads=num_heads,
+            attn_mode=attn_mode,
+            window_size=window_size,
+            shift_window=shift_window,
+            qkv_bias=qkv_bias,
+            use_rope=use_rope,
+            rope_freq=rope_freq,
+            qk_rms_norm=qk_rms_norm,
+        )
+        self.mlp = FeedForwardNet(
+            channels,
+            mlp_ratio=mlp_ratio,
+        )
+
+    def _forward(self, x: torch.Tensor, phases: Optional[torch.Tensor] = None) -> torch.Tensor:
+        h = self.norm1(x)
+        h = self.attn(h, phases=phases)
+        x = x + h
+        h = self.norm2(x)
+        h = self.mlp(h)
+        x = x + h
+        return x
+
+    def forward(self, x: torch.Tensor, phases: Optional[torch.Tensor] = None) -> torch.Tensor:
+        if self.use_checkpoint:
+            return torch.utils.checkpoint.checkpoint(self._forward, x, phases, use_reentrant=False)
+        else:
+            return self._forward(x, phases)
+
+
+class TransformerCrossBlock(nn.Module):
+    """
+    Transformer cross-attention block (MSA + MCA + FFN).
+    """
+    def __init__(
+        self,
+        channels: int,
+        ctx_channels: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        attn_mode: Literal["full", "windowed"] = "full",
+        window_size: Optional[int] = None,
+        shift_window: Optional[Tuple[int, int, int]] = None,
+        use_checkpoint: bool = False,
+        use_rope: bool = False,
+        rope_freq: Tuple[int, int] = (1.0, 10000.0), 
+        qk_rms_norm: bool = False,
+        qk_rms_norm_cross: bool = False,
+        qkv_bias: bool = True,
+        ln_affine: bool = False,
+    ):
+        super().__init__()
+        self.use_checkpoint = use_checkpoint
+        self.norm1 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6)
+        self.norm2 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6)
+        self.norm3 = LayerNorm32(channels, elementwise_affine=ln_affine, eps=1e-6)
+        self.self_attn = MultiHeadAttention(
+            channels,
+            num_heads=num_heads,
+            type="self",
+            attn_mode=attn_mode,
+            window_size=window_size,
+            shift_window=shift_window,
+            qkv_bias=qkv_bias,
+            use_rope=use_rope,
+            rope_freq=rope_freq,
+            qk_rms_norm=qk_rms_norm,
+        )
+        self.cross_attn = MultiHeadAttention(
+            channels,
+            ctx_channels=ctx_channels,
+            num_heads=num_heads,
+            type="cross",
+            attn_mode="full",
+            qkv_bias=qkv_bias,
+            qk_rms_norm=qk_rms_norm_cross,
+        )
+        self.mlp = FeedForwardNet(
+            channels,
+            mlp_ratio=mlp_ratio,
+        )
+
+    def _forward(self, x: torch.Tensor, context: torch.Tensor, phases: Optional[torch.Tensor] = None) -> torch.Tensor:
+        h = self.norm1(x)
+        h = self.self_attn(h, phases=phases)
+        x = x + h
+        h = self.norm2(x)
+        h = self.cross_attn(h, context)
+        x = x + h
+        h = self.norm3(x)
+        h = self.mlp(h)
+        x = x + h
+        return x
+
+    def forward(self, x: torch.Tensor, context: torch.Tensor, phases: Optional[torch.Tensor] = None) -> torch.Tensor:
+        if self.use_checkpoint:
+            return torch.utils.checkpoint.checkpoint(self._forward, x, context, phases, use_reentrant=False)
+        else:
+            return self._forward(x, context, phases)
+        
\ No newline at end of file
diff --git a/pixal3d/modules/transformer/modulated.py b/pixal3d/modules/transformer/modulated.py
new file mode 100644
index 0000000..4444b20
--- /dev/null
+++ b/pixal3d/modules/transformer/modulated.py
@@ -0,0 +1,205 @@
+from typing import *
+import torch
+import torch.nn as nn
+from ..attention import MultiHeadAttention, ProjectAttention, GatedProjectAttention
+from ..norm import LayerNorm32
+from .blocks import FeedForwardNet
+
+
+class ModulatedTransformerBlock(nn.Module):
+    """
+    Transformer block (MSA + FFN) with adaptive layer norm conditioning.
+    """
+    def __init__(
+        self,
+        channels: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        attn_mode: Literal["full", "windowed"] = "full",
+        window_size: Optional[int] = None,
+        shift_window: Optional[Tuple[int, int, int]] = None,
+        use_checkpoint: bool = False,
+        use_rope: bool = False,
+        rope_freq: Tuple[int, int] = (1.0, 10000.0), 
+        qk_rms_norm: bool = False,
+        qkv_bias: bool = True,
+        share_mod: bool = False,
+    ):
+        super().__init__()
+        self.use_checkpoint = use_checkpoint
+        self.share_mod = share_mod
+        self.norm1 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)
+        self.norm2 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)
+        self.attn = MultiHeadAttention(
+            channels,
+            num_heads=num_heads,
+            attn_mode=attn_mode,
+            window_size=window_size,
+            shift_window=shift_window,
+            qkv_bias=qkv_bias,
+            use_rope=use_rope,
+            rope_freq=rope_freq,
+            qk_rms_norm=qk_rms_norm,
+        )
+        self.mlp = FeedForwardNet(
+            channels,
+            mlp_ratio=mlp_ratio,
+        )
+        if not share_mod:
+            self.adaLN_modulation = nn.Sequential(
+                nn.SiLU(),
+                nn.Linear(channels, 6 * channels, bias=True)
+            )
+        else:
+            self.modulation = nn.Parameter(torch.randn(6 * channels) / channels ** 0.5)
+
+    def _forward(self, x: torch.Tensor, mod: torch.Tensor, phases: Optional[torch.Tensor] = None) -> torch.Tensor:
+        if self.share_mod:
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (self.modulation + mod).type(mod.dtype).chunk(6, dim=1)
+        else:
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(mod).chunk(6, dim=1)
+        h = self.norm1(x)
+        h = h * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1)
+        h = self.attn(h, phases=phases)
+        h = h * gate_msa.unsqueeze(1)
+        x = x + h
+        h = self.norm2(x)
+        h = h * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1)
+        h = self.mlp(h)
+        h = h * gate_mlp.unsqueeze(1)
+        x = x + h
+        return x
+
+    def forward(self, x: torch.Tensor, mod: torch.Tensor, phases: Optional[torch.Tensor] = None) -> torch.Tensor:
+        if self.use_checkpoint:
+            return torch.utils.checkpoint.checkpoint(self._forward, x, mod, phases, use_reentrant=False)
+        else:
+            return self._forward(x, mod, phases)
+
+
+class ModulatedTransformerCrossBlock(nn.Module):
+    """
+    Transformer cross-attention block (MSA + MCA + FFN) with adaptive layer norm conditioning.
+    
+    Supports two image attention modes:
+    - "cross": Standard cross-attention with image features
+    - "proj": Projection-based attention with view-aligned features
+    """
+    def __init__(
+        self,
+        channels: int,
+        ctx_channels: int,
+        num_heads: int,
+        mlp_ratio: float = 4.0,
+        attn_mode: Literal["full", "windowed"] = "full",
+        window_size: Optional[int] = None,
+        shift_window: Optional[Tuple[int, int, int]] = None,
+        use_checkpoint: bool = False,
+        use_rope: bool = False,
+        rope_freq: Tuple[int, int] = (1.0, 10000.0), 
+        qk_rms_norm: bool = False,
+        qk_rms_norm_cross: bool = False,
+        qkv_bias: bool = True,
+        share_mod: bool = False,
+        image_attn_mode: Literal["cross", "proj", "gated_proj"] = "cross",
+        proj_in_channels: Optional[int] = None,
+        vae_in_channels: Optional[int] = None,
+    ):
+        super().__init__()
+        self.use_checkpoint = use_checkpoint
+        self.share_mod = share_mod
+        self.image_attn_mode = image_attn_mode
+        
+        self.norm1 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)
+        self.norm2 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
+        self.norm3 = LayerNorm32(channels, elementwise_affine=False, eps=1e-6)
+        self.self_attn = MultiHeadAttention(
+            channels,
+            num_heads=num_heads,
+            type="self",
+            attn_mode=attn_mode,
+            window_size=window_size,
+            shift_window=shift_window,
+            qkv_bias=qkv_bias,
+            use_rope=use_rope,
+            rope_freq=rope_freq,
+            qk_rms_norm=qk_rms_norm,
+        )
+        
+        # Build cross attention based on mode
+        if image_attn_mode == "cross":
+            self.cross_attn = MultiHeadAttention(
+                channels,
+                ctx_channels=ctx_channels,
+                num_heads=num_heads,
+                type="cross",
+                attn_mode="full",
+                qkv_bias=qkv_bias,
+                qk_rms_norm=qk_rms_norm_cross,
+            )
+        elif image_attn_mode == "proj":
+            _proj_in = proj_in_channels if proj_in_channels is not None else ctx_channels
+            cross_attn_block = MultiHeadAttention(
+                channels,
+                ctx_channels=ctx_channels,
+                num_heads=num_heads,
+                type="cross",
+                attn_mode="full",
+                qkv_bias=qkv_bias,
+                qk_rms_norm=qk_rms_norm_cross,
+            )
+            self.cross_attn = ProjectAttention(cross_attn_block, channels, _proj_in)
+        elif image_attn_mode == "gated_proj":
+            _dino_in = proj_in_channels if proj_in_channels is not None else ctx_channels
+            _vae_in = vae_in_channels if vae_in_channels is not None else 16
+            cross_attn_block = MultiHeadAttention(
+                channels,
+                ctx_channels=ctx_channels,
+                num_heads=num_heads,
+                type="cross",
+                attn_mode="full",
+                qkv_bias=qkv_bias,
+                qk_rms_norm=qk_rms_norm_cross,
+            )
+            self.cross_attn = GatedProjectAttention(cross_attn_block, channels, _dino_in, _vae_in)
+        else:
+            raise ValueError(f"Unknown image attention mode: {image_attn_mode}")
+            
+        self.mlp = FeedForwardNet(
+            channels,
+            mlp_ratio=mlp_ratio,
+        )
+        if not share_mod:
+            self.adaLN_modulation = nn.Sequential(
+                nn.SiLU(),
+                nn.Linear(channels, 6 * channels, bias=True)
+            )
+        else:
+            self.modulation = nn.Parameter(torch.randn(6 * channels) / channels ** 0.5)
+
+    def _forward(self, x: torch.Tensor, mod: torch.Tensor, context: torch.Tensor, phases: Optional[torch.Tensor] = None) -> torch.Tensor:
+        if self.share_mod:
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (self.modulation + mod).type(mod.dtype).chunk(6, dim=1)
+        else:
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(mod).chunk(6, dim=1)
+        h = self.norm1(x)
+        h = h * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1)
+        h = self.self_attn(h, phases=phases)
+        h = h * gate_msa.unsqueeze(1)
+        x = x + h
+        h = self.norm2(x)
+        h = self.cross_attn(h, context)
+        x = x + h
+        h = self.norm3(x)
+        h = h * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1)
+        h = self.mlp(h)
+        h = h * gate_mlp.unsqueeze(1)
+        x = x + h
+        return x
+
+    def forward(self, x: torch.Tensor, mod: torch.Tensor, context: torch.Tensor, phases: Optional[torch.Tensor] = None) -> torch.Tensor:
+        if self.use_checkpoint:
+            return torch.utils.checkpoint.checkpoint(self._forward, x, mod, context, phases, use_reentrant=False)
+        else:
+            return self._forward(x, mod, context, phases)
+        
\ No newline at end of file
diff --git a/pixal3d/modules/utils.py b/pixal3d/modules/utils.py
new file mode 100644
index 0000000..5d92d7d
--- /dev/null
+++ b/pixal3d/modules/utils.py
@@ -0,0 +1,87 @@
+import torch
+import torch.nn as nn
+from ..modules import sparse as sp
+
+MIX_PRECISION_MODULES = (
+    nn.Conv1d,
+    nn.Conv2d,
+    nn.Conv3d,
+    nn.ConvTranspose1d,
+    nn.ConvTranspose2d,
+    nn.ConvTranspose3d,
+    nn.Linear,
+    sp.SparseConv3d,
+    sp.SparseInverseConv3d,
+    sp.SparseLinear,
+)
+
+
+def convert_module_to_f16(l):
+    """
+    Convert primitive modules to float16.
+    """
+    if isinstance(l, MIX_PRECISION_MODULES):
+        for p in l.parameters():
+            p.data = p.data.half()
+
+
+def convert_module_to_f32(l):
+    """
+    Convert primitive modules to float32, undoing convert_module_to_f16().
+    """
+    if isinstance(l, MIX_PRECISION_MODULES):
+        for p in l.parameters():
+            p.data = p.data.float()
+
+
+def convert_module_to(l, dtype):
+    """
+    Convert primitive modules to the given dtype.
+    """
+    if isinstance(l, MIX_PRECISION_MODULES):
+        for p in l.parameters():
+            p.data = p.data.to(dtype)
+
+
+def zero_module(module):
+    """
+    Zero out the parameters of a module and return it.
+    """
+    for p in module.parameters():
+        p.detach().zero_()
+    return module
+
+
+def scale_module(module, scale):
+    """
+    Scale the parameters of a module and return it.
+    """
+    for p in module.parameters():
+        p.detach().mul_(scale)
+    return module
+
+
+def modulate(x, shift, scale):
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+def manual_cast(tensor, dtype):
+    """
+    Cast if autocast is not enabled.
+    """
+    if not torch.is_autocast_enabled():
+        return tensor.type(dtype)
+    return tensor
+
+
+def str_to_dtype(dtype_str: str):
+    return {
+        'f16': torch.float16,
+        'fp16': torch.float16,
+        'float16': torch.float16,
+        'bf16': torch.bfloat16,
+        'bfloat16': torch.bfloat16,
+        'f32': torch.float32,
+        'fp32': torch.float32,
+        'float32': torch.float32,
+    }[dtype_str]
diff --git a/pixal3d/pipelines/__init__.py b/pixal3d/pipelines/__init__.py
new file mode 100644
index 0000000..3bc1dfa
--- /dev/null
+++ b/pixal3d/pipelines/__init__.py
@@ -0,0 +1,55 @@
+import importlib
+
+__attributes = {
+    "Trellis2ImageTo3DPipeline": "trellis2_image_to_3d",
+    "Trellis2TexturingPipeline": "trellis2_texturing",
+    "Pixal3DImageTo3DPipeline": "pixal3d_image_to_3d",
+}
+
+
+__submodules = ['samplers', 'rembg']
+
+__all__ = list(__attributes.keys()) + __submodules
+
+def __getattr__(name):
+    if name not in globals():
+        if name in __attributes:
+            module_name = __attributes[name]
+            module = importlib.import_module(f".{module_name}", __name__)
+            globals()[name] = getattr(module, name)
+        elif name in __submodules:
+            module = importlib.import_module(f".{name}", __name__)
+            globals()[name] = module
+        else:
+            raise AttributeError(f"module {__name__} has no attribute {name}")
+    return globals()[name]
+
+
+def from_pretrained(path: str):
+    """
+    Load a pipeline from a model folder or a Hugging Face model hub.
+
+    Args:
+        path: The path to the model. Can be either local path or a Hugging Face model name.
+    """
+    import os
+    import json
+    is_local = os.path.exists(f"{path}/pipeline.json")
+
+    if is_local:
+        config_file = f"{path}/pipeline.json"
+    else:
+        from huggingface_hub import hf_hub_download
+        config_file = hf_hub_download(path, "pipeline.json")
+
+    with open(config_file, 'r') as f:
+        config = json.load(f)
+    return globals()[config['name']].from_pretrained(path)
+
+
+# For PyLance
+if __name__ == '__main__':
+    from . import samplers, rembg
+    from .trellis2_image_to_3d import Trellis2ImageTo3DPipeline
+    from .trellis2_texturing import Trellis2TexturingPipeline
+    from .pixal3d_image_to_3d import Pixal3DImageTo3DPipeline
diff --git a/pixal3d/pipelines/base.py b/pixal3d/pipelines/base.py
new file mode 100644
index 0000000..331e1ed
--- /dev/null
+++ b/pixal3d/pipelines/base.py
@@ -0,0 +1,72 @@
+from typing import *
+import torch
+import torch.nn as nn
+from .. import models
+
+
+class Pipeline:
+    """
+    A base class for pipelines.
+    """
+    def __init__(
+        self,
+        models: dict[str, nn.Module] = None,
+    ):
+        if models is None:
+            return
+        self.models = models
+        for model in self.models.values():
+            model.eval()
+
+    @classmethod
+    def from_pretrained(cls, path: str, config_file: str = "pipeline.json") -> "Pipeline":
+        """
+        Load a pretrained model.
+        """
+        import os
+        import json
+        is_local = os.path.exists(f"{path}/{config_file}")
+
+        if is_local:
+            config_file = f"{path}/{config_file}"
+        else:
+            from huggingface_hub import hf_hub_download
+            config_file = hf_hub_download(path, config_file)
+
+        with open(config_file, 'r') as f:
+            args = json.load(f)['args']
+
+        _models = {}
+        for k, v in args['models'].items():
+            if hasattr(cls, 'model_names_to_load') and k not in cls.model_names_to_load:
+                continue
+            try:
+                _models[k] = models.from_pretrained(f"{path}/{v}")
+            except Exception as e:
+                _models[k] = models.from_pretrained(v)
+
+        new_pipeline = cls(_models)
+        new_pipeline._pretrained_args = args
+        return new_pipeline
+
+    @property
+    def device(self) -> torch.device:
+        if hasattr(self, '_device'):
+            return self._device
+        for model in self.models.values():
+            if hasattr(model, 'device'):
+                return model.device
+        for model in self.models.values():
+            if hasattr(model, 'parameters'):
+                return next(model.parameters()).device
+        raise RuntimeError("No device found.")
+
+    def to(self, device: torch.device) -> None:
+        for model in self.models.values():
+            model.to(device)
+
+    def cuda(self) -> None:
+        self.to(torch.device("cuda"))
+
+    def cpu(self) -> None:
+        self.to(torch.device("cpu"))
\ No newline at end of file
diff --git a/pixal3d/pipelines/pixal3d_image_to_3d.py b/pixal3d/pipelines/pixal3d_image_to_3d.py
new file mode 100644
index 0000000..7129482
--- /dev/null
+++ b/pixal3d/pipelines/pixal3d_image_to_3d.py
@@ -0,0 +1,783 @@
+from typing import *
+import torch
+import torch.nn as nn
+import numpy as np
+from PIL import Image
+from .base import Pipeline
+from . import samplers, rembg
+from ..modules.sparse import SparseTensor
+from ..modules import image_feature_extractor
+from ..representations import Mesh, MeshWithVoxel
+
+
+class Pixal3DImageTo3DPipeline(Pipeline):
+    """
+    Pipeline for inferring Pixal3D (proj mode) image-to-3D models.
+
+    Based on Trellis2 pipeline, using proj mode for inference.
+    Each stage (SS, Shape 512, Shape 1024, Tex 1024) has its own image_cond_model (DinoV3ProjFeatureExtractor).
+    Condition building uses camera-aware projection (requires camera_angle_x, distance, mesh_scale parameters).
+
+    Args:
+        models (dict[str, nn.Module]): The models to use in the pipeline.
+        sparse_structure_sampler (samplers.Sampler): The sampler for the sparse structure.
+        shape_slat_sampler (samplers.Sampler): The sampler for the structured latent.
+        tex_slat_sampler (samplers.Sampler): The sampler for the texture latent.
+        sparse_structure_sampler_params (dict): The parameters for the sparse structure sampler.
+        shape_slat_sampler_params (dict): The parameters for the structured latent sampler.
+        tex_slat_sampler_params (dict): The parameters for the texture latent sampler.
+        shape_slat_normalization (dict): The normalization parameters for the structured latent.
+        tex_slat_normalization (dict): The normalization parameters for the texture latent.
+        image_cond_model_ss (nn.Module): Proj image cond model for sparse structure stage.
+        image_cond_model_shape_512 (nn.Module): Proj image cond model for shape LR (512) stage.
+        image_cond_model_shape_1024 (nn.Module): Proj image cond model for shape HR (1024) stage.
+        image_cond_model_tex_1024 (nn.Module): Proj image cond model for texture (1024) stage.
+        rembg_model (Callable): The model for removing background.
+        low_vram (bool): Whether to use low-VRAM mode.
+    """
+    model_names_to_load = [
+        'sparse_structure_flow_model',
+        'sparse_structure_decoder',
+        'shape_slat_flow_model_512',
+        'shape_slat_flow_model_1024',
+        'shape_slat_decoder',
+        'tex_slat_flow_model_512',
+        'tex_slat_flow_model_1024',
+        'tex_slat_decoder',
+    ]
+
+    def __init__(
+        self,
+        models: dict[str, nn.Module] = None,
+        sparse_structure_sampler: samplers.Sampler = None,
+        shape_slat_sampler: samplers.Sampler = None,
+        tex_slat_sampler: samplers.Sampler = None,
+        sparse_structure_sampler_params: dict = None,
+        shape_slat_sampler_params: dict = None,
+        tex_slat_sampler_params: dict = None,
+        shape_slat_normalization: dict = None,
+        tex_slat_normalization: dict = None,
+        image_cond_model_ss: nn.Module = None,
+        image_cond_model_shape_512: nn.Module = None,
+        image_cond_model_shape_1024: nn.Module = None,
+        image_cond_model_tex_1024: nn.Module = None,
+        rembg_model: Callable = None,
+        low_vram: bool = True,
+        default_pipeline_type: str = '1024_cascade',
+    ):
+        if models is None:
+            return
+        super().__init__(models)
+        self.sparse_structure_sampler = sparse_structure_sampler
+        self.shape_slat_sampler = shape_slat_sampler
+        self.tex_slat_sampler = tex_slat_sampler
+        self.sparse_structure_sampler_params = sparse_structure_sampler_params
+        self.shape_slat_sampler_params = shape_slat_sampler_params
+        self.tex_slat_sampler_params = tex_slat_sampler_params
+        self.shape_slat_normalization = shape_slat_normalization
+        self.tex_slat_normalization = tex_slat_normalization
+        self.image_cond_model_ss = image_cond_model_ss
+        self.image_cond_model_shape_512 = image_cond_model_shape_512
+        self.image_cond_model_shape_1024 = image_cond_model_shape_1024
+        self.image_cond_model_tex_1024 = image_cond_model_tex_1024
+        self.rembg_model = rembg_model
+        self.low_vram = low_vram
+        self.default_pipeline_type = default_pipeline_type
+        self.pbr_attr_layout = {
+            'base_color': slice(0, 3),
+            'metallic': slice(3, 4),
+            'roughness': slice(4, 5),
+            'alpha': slice(5, 6),
+        }
+        self._device = 'cpu'
+
+    @classmethod
+    def from_pretrained(cls, path: str, config_file: str = "pipeline.json") -> "Pixal3DImageTo3DPipeline":
+        """
+        Load a pretrained model.
+
+        Args:
+            path (str): The path to the model. Can be either local path or a Hugging Face repository.
+        """
+        pipeline = super().from_pretrained(path, config_file)
+        args = pipeline._pretrained_args
+
+        pipeline.sparse_structure_sampler = getattr(samplers, args['sparse_structure_sampler']['name'])(**args['sparse_structure_sampler']['args'])
+        pipeline.sparse_structure_sampler_params = args['sparse_structure_sampler']['params']
+
+        pipeline.shape_slat_sampler = getattr(samplers, args['shape_slat_sampler']['name'])(**args['shape_slat_sampler']['args'])
+        pipeline.shape_slat_sampler_params = args['shape_slat_sampler']['params']
+
+        pipeline.tex_slat_sampler = getattr(samplers, args['tex_slat_sampler']['name'])(**args['tex_slat_sampler']['args'])
+        pipeline.tex_slat_sampler_params = args['tex_slat_sampler']['params']
+
+        pipeline.shape_slat_normalization = args['shape_slat_normalization']
+        pipeline.tex_slat_normalization = args['tex_slat_normalization']
+
+        # Proj mode: image_cond_models need to be loaded externally, set to None here
+        pipeline.image_cond_model_ss = None
+        pipeline.image_cond_model_shape_512 = None
+        pipeline.image_cond_model_shape_1024 = None
+        pipeline.image_cond_model_tex_1024 = None
+
+        pipeline.rembg_model = getattr(rembg, args['rembg_model']['name'])(**args['rembg_model']['args'])
+        
+        pipeline.low_vram = args.get('low_vram', True)
+        pipeline.default_pipeline_type = args.get('default_pipeline_type', '1024_cascade')
+        pipeline.pbr_attr_layout = {
+            'base_color': slice(0, 3),
+            'metallic': slice(3, 4),
+            'roughness': slice(4, 5),
+            'alpha': slice(5, 6),
+        }
+        pipeline._device = 'cpu'
+
+        return pipeline
+
+    def to(self, device: torch.device) -> None:
+        self._device = device
+        if not self.low_vram:
+            super().to(device)
+            if self.rembg_model is not None:
+                self.rembg_model.to(device)
+
+    def preprocess_image(self, input: Image.Image, bg_color: tuple = (0, 0, 0)) -> Image.Image:
+        """
+        Preprocess the input image.
+
+        Args:
+            input: Input image (RGB or RGBA).
+            bg_color: Background color (R, G, B) in 0~255. Default black (0,0,0).
+        """
+        # if has alpha channel, use it directly; otherwise, remove background
+        has_alpha = False
+        if input.mode == 'RGBA':
+            alpha = np.array(input)[:, :, 3]
+            if not np.all(alpha == 255):
+                has_alpha = True
+        max_size = max(input.size)
+        scale = min(1, 1024 / max_size)
+        if scale < 1:
+            input = input.resize((int(input.width * scale), int(input.height * scale)), Image.Resampling.LANCZOS)
+        if has_alpha:
+            output = input
+        else:
+            input = input.convert('RGB')
+            if self.low_vram:
+                self.rembg_model.to(self.device)
+            output = self.rembg_model(input)
+            if self.low_vram:
+                self.rembg_model.cpu()
+        output_np = np.array(output)
+        alpha = output_np[:, :, 3]
+        bbox = np.argwhere(alpha > 0.8 * 255)
+        bbox = np.min(bbox[:, 1]), np.min(bbox[:, 0]), np.max(bbox[:, 1]), np.max(bbox[:, 0])
+        center = (bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2
+        size = max(bbox[2] - bbox[0], bbox[3] - bbox[1])
+        size = int(size * 1.1)
+        bbox = center[0] - size // 2, center[1] - size // 2, center[0] + size // 2, center[1] + size // 2
+        output = output.crop(bbox)  # type: ignore
+        output = np.array(output).astype(np.float32) / 255
+        rgb = output[:, :, :3]
+        a = output[:, :, 3:4]
+        bg = np.array(bg_color, dtype=np.float32) / 255.0
+        output = rgb * a + bg * (1.0 - a)
+        output = Image.fromarray((np.clip(output, 0, 1) * 255).astype(np.uint8))
+        return output
+
+    # =========================================================================
+    # Proj mode condition building
+    # =========================================================================
+
+    @torch.no_grad()
+    def get_proj_cond_ss(
+        self,
+        image: list,
+        camera_angle_x: float = 0.8575560450553894,
+        distance: float = 2.0,
+        mesh_scale: float = 1.0,
+    ) -> dict:
+        """
+        Get proj conditioning for sparse structure stage.
+
+        Args:
+            image: List of PIL images.
+            camera_angle_x: Camera horizontal FOV in radians.
+            distance: Camera distance.
+            mesh_scale: Mesh scale.
+
+        Returns:
+            dict with 'cond' and 'neg_cond', each containing {'global': ..., 'proj': ...}
+        """
+        device = self.device
+        image_cond_model = self.image_cond_model_ss
+        if self.low_vram:
+            image_cond_model.to(device)
+        cam_angle = torch.tensor([camera_angle_x], device=device)
+        dist_tensor = torch.tensor([distance], device=device)
+        scale_tensor = torch.tensor([mesh_scale], device=device)
+        z_global, z_proj = image_cond_model(
+            image, camera_angle_x=cam_angle, distance=dist_tensor, mesh_scale=scale_tensor,
+        )
+        if self.low_vram:
+            image_cond_model.cpu()
+        return {
+            'cond': {'global': z_global, 'proj': z_proj},
+            'neg_cond': {'global': torch.zeros_like(z_global), 'proj': torch.zeros_like(z_proj)},
+        }
+
+    @torch.no_grad()
+    def get_proj_cond_shape(
+        self,
+        image_cond_model: nn.Module,
+        image: list,
+        coords: torch.Tensor,
+        camera_angle_x: float = 0.8575560450553894,
+        distance: float = 2.0,
+        mesh_scale: float = 1.0,
+        grid_resolution_override: int = None,
+    ) -> dict:
+        """
+        Get proj conditioning for shape/texture stages (sparse-token aligned).
+
+        Args:
+            image_cond_model: The proj image cond model for this stage.
+            image: List of PIL images.
+            coords: Sparse structure coordinates [N, 4] (batch_idx, x, y, z).
+            camera_angle_x: Camera horizontal FOV in radians.
+            distance: Camera distance.
+            mesh_scale: Mesh scale.
+            grid_resolution_override: Override the grid resolution if not None.
+
+        Returns:
+            dict with 'cond' and 'neg_cond', each containing {'global': ..., 'proj': SparseTensor}
+        """
+        device = self.device
+        if self.low_vram:
+            image_cond_model.to(device)
+
+        orig_grid_res = image_cond_model.grid_resolution
+        if grid_resolution_override is not None and grid_resolution_override != orig_grid_res:
+            image_cond_model.grid_resolution = grid_resolution_override
+            image_cond_model.proj_grid = image_cond_model.proj_grid.__class__(
+                grid_resolution=grid_resolution_override,
+                image_resolution=image_cond_model.proj_grid.image_resolution,
+            ).to(device)
+
+        B = 1
+        cam_angle = torch.tensor([camera_angle_x], device=device)
+        dist_tensor = torch.tensor([distance], device=device)
+        scale_tensor = torch.tensor([mesh_scale], device=device)
+        z_global, z_proj = image_cond_model(
+            image, camera_angle_x=cam_angle, distance=dist_tensor, mesh_scale=scale_tensor,
+        )
+        grid_res = image_cond_model.grid_resolution
+        z_proj_grid = z_proj.reshape(B, grid_res, grid_res, grid_res, -1)
+        batch_indices = coords[:, 0].long()
+        x_coords = coords[:, 1].long()
+        y_coords = coords[:, 2].long()
+        z_coords = coords[:, 3].long()
+        z_proj_sparse = z_proj_grid[batch_indices, x_coords, y_coords, z_coords]
+        z_proj_st = SparseTensor(feats=z_proj_sparse, coords=coords)
+
+        if grid_resolution_override is not None and grid_resolution_override != orig_grid_res:
+            image_cond_model.grid_resolution = orig_grid_res
+            image_cond_model.proj_grid = image_cond_model.proj_grid.__class__(
+                grid_resolution=orig_grid_res,
+                image_resolution=image_cond_model.proj_grid.image_resolution,
+            ).to(device)
+
+        if self.low_vram:
+            image_cond_model.cpu()
+        return {
+            'cond': {'global': z_global, 'proj': z_proj_st},
+            'neg_cond': {'global': torch.zeros_like(z_global), 'proj': SparseTensor(feats=torch.zeros_like(z_proj_sparse), coords=coords)},
+        }
+
+    # =========================================================================
+    # Sampling methods (consistent with Trellis2)
+    # =========================================================================
+
+    def sample_sparse_structure(
+        self,
+        cond: dict,
+        resolution: int,
+        num_samples: int = 1,
+        sampler_params: dict = {},
+    ) -> torch.Tensor:
+        """
+        Sample sparse structures with the given conditioning.
+        
+        Args:
+            cond (dict): The conditioning information.
+            resolution (int): The resolution of the sparse structure.
+            num_samples (int): The number of samples to generate.
+            sampler_params (dict): Additional parameters for the sampler.
+        """
+        # Sample sparse structure latent
+        flow_model = self.models['sparse_structure_flow_model']
+        reso = flow_model.resolution
+        in_channels = flow_model.in_channels
+        noise = torch.randn(num_samples, in_channels, reso, reso, reso).to(self.device)
+        sampler_params = {**self.sparse_structure_sampler_params, **sampler_params}
+        if self.low_vram:
+            flow_model.to(self.device)
+        z_s = self.sparse_structure_sampler.sample(
+            flow_model,
+            noise,
+            **cond,
+            **sampler_params,
+            verbose=True,
+            tqdm_desc="Sampling sparse structure (proj)",
+        ).samples
+        if self.low_vram:
+            flow_model.cpu()
+        
+        # Decode sparse structure latent
+        decoder = self.models['sparse_structure_decoder']
+        if self.low_vram:
+            decoder.to(self.device)
+        decoded = decoder(z_s)>0
+        if self.low_vram:
+            decoder.cpu()
+        if resolution != decoded.shape[2]:
+            ratio = decoded.shape[2] // resolution
+            decoded = torch.nn.functional.max_pool3d(decoded.float(), ratio, ratio, 0) > 0.5
+        coords = torch.argwhere(decoded)[:, [0, 2, 3, 4]].int()
+
+        return coords
+
+    def sample_shape_slat(
+        self,
+        cond: dict,
+        flow_model,
+        coords: torch.Tensor,
+        sampler_params: dict = {},
+    ) -> SparseTensor:
+        """
+        Sample structured latent with the given conditioning.
+        
+        Args:
+            cond (dict): The conditioning information.
+            coords (torch.Tensor): The coordinates of the sparse structure.
+            sampler_params (dict): Additional parameters for the sampler.
+        """
+        # Sample structured latent
+        noise = SparseTensor(
+            feats=torch.randn(coords.shape[0], flow_model.in_channels).to(self.device),
+            coords=coords,
+        )
+        sampler_params = {**self.shape_slat_sampler_params, **sampler_params}
+        if self.low_vram:
+            flow_model.to(self.device)
+        slat = self.shape_slat_sampler.sample(
+            flow_model,
+            noise,
+            **cond,
+            **sampler_params,
+            verbose=True,
+            tqdm_desc="Sampling shape SLat (proj)",
+        ).samples
+        if self.low_vram:
+            flow_model.cpu()
+
+        std = torch.tensor(self.shape_slat_normalization['std'])[None].to(slat.device)
+        mean = torch.tensor(self.shape_slat_normalization['mean'])[None].to(slat.device)
+        slat = slat * std + mean
+        
+        return slat
+    
+    def sample_shape_slat_cascade(
+        self,
+        lr_cond: dict,
+        cond: dict,
+        flow_model_lr,
+        flow_model,
+        lr_resolution: int,
+        resolution: int,
+        coords: torch.Tensor,
+        sampler_params: dict = {},
+        max_num_tokens: int = 49152,
+    ) -> SparseTensor:
+        """
+        Sample structured latent with cascade (LR → HR).
+        
+        Args:
+            lr_cond (dict): The conditioning information for LR stage.
+            cond (dict): The conditioning information for HR stage.
+            flow_model_lr: LR flow model.
+            flow_model: HR flow model.
+            lr_resolution (int): LR resolution.
+            resolution (int): Target HR resolution.
+            coords (torch.Tensor): The coordinates of the sparse structure.
+            sampler_params (dict): Additional parameters for the sampler.
+            max_num_tokens (int): Maximum number of tokens.
+        """
+        # LR
+        noise = SparseTensor(
+            feats=torch.randn(coords.shape[0], flow_model_lr.in_channels).to(self.device),
+            coords=coords,
+        )
+        sampler_params = {**self.shape_slat_sampler_params, **sampler_params}
+        if self.low_vram:
+            flow_model_lr.to(self.device)
+        slat = self.shape_slat_sampler.sample(
+            flow_model_lr,
+            noise,
+            **lr_cond,
+            **sampler_params,
+            verbose=True,
+            tqdm_desc="Sampling LR shape SLat (proj, 512)",
+        ).samples
+        if self.low_vram:
+            flow_model_lr.cpu()
+        std = torch.tensor(self.shape_slat_normalization['std'])[None].to(slat.device)
+        mean = torch.tensor(self.shape_slat_normalization['mean'])[None].to(slat.device)
+        slat = slat * std + mean
+        
+        # Upsample
+        if self.low_vram:
+            self.models['shape_slat_decoder'].to(self.device)
+            self.models['shape_slat_decoder'].low_vram = True
+        hr_coords = self.models['shape_slat_decoder'].upsample(slat, upsample_times=4)
+        if self.low_vram:
+            self.models['shape_slat_decoder'].cpu()
+            self.models['shape_slat_decoder'].low_vram = False
+        hr_resolution = resolution
+        while True:
+            quant_coords = torch.cat([
+                hr_coords[:, :1],
+                ((hr_coords[:, 1:] + 0.5) / lr_resolution * (hr_resolution // 16)).int(),
+            ], dim=1)
+            coords = quant_coords.unique(dim=0)
+            num_tokens = coords.shape[0]
+            if num_tokens < max_num_tokens or hr_resolution == 1024:
+                if hr_resolution != resolution:
+                    print(f"Due to the limited number of tokens, the resolution is reduced to {hr_resolution}.")
+                break
+            hr_resolution -= 128
+        
+        # Sample structured latent (HR)
+        noise = SparseTensor(
+            feats=torch.randn(coords.shape[0], flow_model.in_channels).to(self.device),
+            coords=coords,
+        )
+        sampler_params = {**self.shape_slat_sampler_params, **sampler_params}
+        if self.low_vram:
+            flow_model.to(self.device)
+        slat = self.shape_slat_sampler.sample(
+            flow_model,
+            noise,
+            **cond,
+            **sampler_params,
+            verbose=True,
+            tqdm_desc=f"Sampling HR shape SLat (proj, {hr_resolution})",
+        ).samples
+        if self.low_vram:
+            flow_model.cpu()
+
+        std = torch.tensor(self.shape_slat_normalization['std'])[None].to(slat.device)
+        mean = torch.tensor(self.shape_slat_normalization['mean'])[None].to(slat.device)
+        slat = slat * std + mean
+        
+        return slat, hr_resolution
+
+    def decode_shape_slat(
+        self,
+        slat: SparseTensor,
+        resolution: int,
+    ) -> Tuple[List[Mesh], List[SparseTensor]]:
+        """
+        Decode the structured latent.
+
+        Args:
+            slat (SparseTensor): The structured latent.
+
+        Returns:
+            List[Mesh]: The decoded meshes.
+            List[SparseTensor]: The decoded substructures.
+        """
+        self.models['shape_slat_decoder'].set_resolution(resolution)
+        if self.low_vram:
+            self.models['shape_slat_decoder'].to(self.device)
+            self.models['shape_slat_decoder'].low_vram = True
+        ret = self.models['shape_slat_decoder'](slat, return_subs=True)
+        if self.low_vram:
+            self.models['shape_slat_decoder'].cpu()
+            self.models['shape_slat_decoder'].low_vram = False
+        return ret
+    
+    def sample_tex_slat(
+        self,
+        cond: dict,
+        flow_model,
+        shape_slat: SparseTensor,
+        sampler_params: dict = {},
+    ) -> SparseTensor:
+        """
+        Sample texture structured latent with the given conditioning.
+        
+        Args:
+            cond (dict): The conditioning information.
+            shape_slat (SparseTensor): The structured latent for shape.
+            sampler_params (dict): Additional parameters for the sampler.
+        """
+        # Sample structured latent
+        std = torch.tensor(self.shape_slat_normalization['std'])[None].to(shape_slat.device)
+        mean = torch.tensor(self.shape_slat_normalization['mean'])[None].to(shape_slat.device)
+        shape_slat = (shape_slat - mean) / std
+
+        in_channels = flow_model.in_channels if isinstance(flow_model, nn.Module) else flow_model[0].in_channels
+        noise = shape_slat.replace(feats=torch.randn(shape_slat.coords.shape[0], in_channels - shape_slat.feats.shape[1]).to(self.device))
+        sampler_params = {**self.tex_slat_sampler_params, **sampler_params}
+        if self.low_vram:
+            flow_model.to(self.device)
+        slat = self.tex_slat_sampler.sample(
+            flow_model,
+            noise,
+            concat_cond=shape_slat,
+            **cond,
+            **sampler_params,
+            verbose=True,
+            tqdm_desc="Sampling texture SLat (proj)",
+        ).samples
+        if self.low_vram:
+            flow_model.cpu()
+
+        std = torch.tensor(self.tex_slat_normalization['std'])[None].to(slat.device)
+        mean = torch.tensor(self.tex_slat_normalization['mean'])[None].to(slat.device)
+        slat = slat * std + mean
+        
+        return slat
+
+    def decode_tex_slat(
+        self,
+        slat: SparseTensor,
+        subs: List[SparseTensor],
+    ) -> SparseTensor:
+        """
+        Decode the structured latent.
+
+        Args:
+            slat (SparseTensor): The structured latent.
+
+        Returns:
+            SparseTensor: The decoded texture voxels
+        """
+        if self.low_vram:
+            self.models['tex_slat_decoder'].to(self.device)
+        ret = self.models['tex_slat_decoder'](slat, guide_subs=subs) * 0.5 + 0.5
+        if self.low_vram:
+            self.models['tex_slat_decoder'].cpu()
+        return ret
+    
+    @torch.no_grad()
+    def decode_latent(
+        self,
+        shape_slat: SparseTensor,
+        tex_slat: SparseTensor,
+        resolution: int,
+    ) -> List[MeshWithVoxel]:
+        """
+        Decode the latent codes.
+
+        Args:
+            shape_slat (SparseTensor): The structured latent for shape.
+            tex_slat (SparseTensor): The structured latent for texture.
+            resolution (int): The resolution of the output.
+        """
+        meshes, subs = self.decode_shape_slat(shape_slat, resolution)
+        tex_voxels = self.decode_tex_slat(tex_slat, subs)
+        out_mesh = []
+        torch.cuda.synchronize()
+        for m, v in zip(meshes, tex_voxels):
+            m.fill_holes()
+            out_mesh.append(
+                MeshWithVoxel(
+                    m.vertices, m.faces,
+                    origin = [-0.5, -0.5, -0.5],
+                    voxel_size = 1 / resolution,
+                    coords = v.coords[:, 1:],
+                    attrs = v.feats,
+                    voxel_shape = torch.Size([*v.shape, *v.spatial_shape]),
+                    layout=self.pbr_attr_layout
+                )
+            )
+        return out_mesh
+    
+    @torch.no_grad()
+    def run(
+        self,
+        image: Image.Image,
+        camera_params: dict,
+        num_samples: int = 1,
+        seed: int = 42,
+        sparse_structure_sampler_params: dict = {},
+        shape_slat_sampler_params: dict = {},
+        tex_slat_sampler_params: dict = {},
+        preprocess_image: bool = True,
+        return_latent: bool = False,
+        pipeline_type: Optional[str] = None,
+        max_num_tokens: int = 49152,
+    ) -> List[MeshWithVoxel]:
+        """
+        Run the Pixal3D pipeline (proj mode, cascade).
+
+        Args:
+            image (Image.Image): The image prompt.
+            camera_params (dict): Camera parameters with keys:
+                - camera_angle_x (float): Horizontal FOV in radians.
+                - distance (float): Camera distance.
+                - mesh_scale (float): Mesh scale factor.
+            num_samples (int): The number of samples to generate.
+            seed (int): The random seed.
+            sparse_structure_sampler_params (dict): Additional parameters for the sparse structure sampler.
+            shape_slat_sampler_params (dict): Additional parameters for the shape SLat sampler.
+            tex_slat_sampler_params (dict): Additional parameters for the texture SLat sampler.
+            preprocess_image (bool): Whether to preprocess the image.
+            return_latent (bool): Whether to return the latent codes.
+            pipeline_type (str): The type of the pipeline. Options: '1024_cascade', '1536_cascade'.
+            max_num_tokens (int): The maximum number of tokens to use.
+        """
+        # Check pipeline type
+        pipeline_type = pipeline_type or self.default_pipeline_type
+        if pipeline_type == '1024_cascade':
+            assert 'shape_slat_flow_model_512' in self.models, "No 512 resolution shape SLat flow model found."
+            assert 'shape_slat_flow_model_1024' in self.models, "No 1024 resolution shape SLat flow model found."
+            assert 'tex_slat_flow_model_1024' in self.models, "No 1024 resolution texture SLat flow model found."
+            hr_resolution = 1024
+        elif pipeline_type == '1536_cascade':
+            assert 'shape_slat_flow_model_512' in self.models, "No 512 resolution shape SLat flow model found."
+            assert 'shape_slat_flow_model_1024' in self.models, "No 1024 resolution shape SLat flow model found."
+            assert 'tex_slat_flow_model_1024' in self.models, "No 1024 resolution texture SLat flow model found."
+            hr_resolution = 1536
+        else:
+            raise ValueError(f"Invalid pipeline type for Pixal3D proj mode: {pipeline_type}. "
+                             f"Supported: '1024_cascade', '1536_cascade'.")
+
+        # Validate image_cond_models are set
+        assert self.image_cond_model_ss is not None, "image_cond_model_ss not set."
+        assert self.image_cond_model_shape_512 is not None, "image_cond_model_shape_512 not set."
+        assert self.image_cond_model_shape_1024 is not None, "image_cond_model_shape_1024 not set."
+        assert self.image_cond_model_tex_1024 is not None, "image_cond_model_tex_1024 not set."
+
+        # Extract camera params
+        camera_angle_x = camera_params['camera_angle_x']
+        distance = camera_params['distance']
+        mesh_scale = camera_params.get('mesh_scale', 1.0)
+        
+        if preprocess_image:
+            image = self.preprocess_image(image)
+        torch.manual_seed(seed)
+
+        # ---- Stage 1: Sparse Structure (proj) ----
+        cond_ss = self.get_proj_cond_ss(
+            [image],
+            camera_angle_x=camera_angle_x,
+            distance=distance,
+            mesh_scale=mesh_scale,
+        )
+        ss_res = 32
+        coords = self.sample_sparse_structure(
+            cond_ss, ss_res,
+            num_samples, sparse_structure_sampler_params
+        )
+        del cond_ss
+        torch.cuda.empty_cache()
+
+        # ---- Stage 2: Shape LR 512 (proj) ----
+        cond_shape_lr = self.get_proj_cond_shape(
+            self.image_cond_model_shape_512, [image], coords,
+            camera_angle_x=camera_angle_x,
+            distance=distance,
+            mesh_scale=mesh_scale,
+        )
+        lr_slat = self.sample_shape_slat(
+            cond_shape_lr, self.models['shape_slat_flow_model_512'],
+            coords, shape_slat_sampler_params
+        )
+        del cond_shape_lr
+        torch.cuda.empty_cache()
+
+        # ---- Stage 3a: Upsample LR → HR ----
+        if self.low_vram:
+            self.models['shape_slat_decoder'].to(self.device)
+            self.models['shape_slat_decoder'].low_vram = True
+        hr_coords = self.models['shape_slat_decoder'].upsample(lr_slat, upsample_times=4)
+        if self.low_vram:
+            self.models['shape_slat_decoder'].cpu()
+            self.models['shape_slat_decoder'].low_vram = False
+
+        lr_resolution = 512
+        actual_hr_resolution = hr_resolution
+        while True:
+            grid_res = actual_hr_resolution // 16
+            quant_coords = torch.cat([
+                hr_coords[:, :1],
+                ((hr_coords[:, 1:] + 0.5) / lr_resolution * (grid_res - 1)).round().int(),
+            ], dim=1)
+            hr_coords_unique = quant_coords.unique(dim=0)
+            num_tokens = hr_coords_unique.shape[0]
+            if num_tokens < max_num_tokens or actual_hr_resolution == 1024:
+                break
+            actual_hr_resolution -= 128
+
+        actual_grid_res = actual_hr_resolution // 16
+        del lr_slat, hr_coords, quant_coords
+        torch.cuda.empty_cache()
+
+        # ---- Stage 3b: Shape HR (proj) ----
+        cond_shape_hr = self.get_proj_cond_shape(
+            self.image_cond_model_shape_1024, [image], hr_coords_unique,
+            camera_angle_x=camera_angle_x,
+            distance=distance,
+            mesh_scale=mesh_scale,
+            grid_resolution_override=actual_grid_res,
+        )
+        noise_hr = SparseTensor(
+            feats=torch.randn(hr_coords_unique.shape[0], self.models['shape_slat_flow_model_1024'].in_channels).to(self.device),
+            coords=hr_coords_unique,
+        )
+        sampler_params_hr = {**self.shape_slat_sampler_params, **shape_slat_sampler_params}
+        flow_model_hr = self.models['shape_slat_flow_model_1024']
+        if self.low_vram:
+            flow_model_hr.to(self.device)
+        hr_slat = self.shape_slat_sampler.sample(
+            flow_model_hr,
+            noise_hr,
+            **cond_shape_hr,
+            **sampler_params_hr,
+            verbose=True,
+            tqdm_desc=f"Sampling HR shape SLat (proj, {actual_hr_resolution})",
+        ).samples
+        if self.low_vram:
+            flow_model_hr.cpu()
+        std = torch.tensor(self.shape_slat_normalization['std'])[None].to(hr_slat.device)
+        mean = torch.tensor(self.shape_slat_normalization['mean'])[None].to(hr_slat.device)
+        shape_slat = hr_slat * std + mean
+        del cond_shape_hr, noise_hr, hr_slat, hr_coords_unique
+        torch.cuda.empty_cache()
+
+        # ---- Stage 4: Texture (proj) ----
+        tex_grid_res = actual_hr_resolution // 16
+        cond_tex = self.get_proj_cond_shape(
+            self.image_cond_model_tex_1024, [image], shape_slat.coords,
+            camera_angle_x=camera_angle_x,
+            distance=distance,
+            mesh_scale=mesh_scale,
+            grid_resolution_override=tex_grid_res,
+        )
+        tex_slat = self.sample_tex_slat(
+            cond_tex, self.models['tex_slat_flow_model_1024'],
+            shape_slat, tex_slat_sampler_params
+        )
+        del cond_tex
+        torch.cuda.empty_cache()
+
+        # ---- Stage 5: Decode ----
+        res = actual_hr_resolution
+        out_mesh = self.decode_latent(shape_slat, tex_slat, res)
+        if return_latent:
+            return out_mesh, (shape_slat, tex_slat, res)
+        else:
+            return out_mesh
diff --git a/pixal3d/pipelines/rembg/BiRefNet.py b/pixal3d/pipelines/rembg/BiRefNet.py
new file mode 100644
index 0000000..c71a992
--- /dev/null
+++ b/pixal3d/pipelines/rembg/BiRefNet.py
@@ -0,0 +1,42 @@
+from typing import *
+from transformers import AutoModelForImageSegmentation
+import torch
+from torchvision import transforms
+from PIL import Image
+
+
+class BiRefNet:
+    def __init__(self, model_name: str = "ZhengPeng7/BiRefNet"):
+        self.model = AutoModelForImageSegmentation.from_pretrained(
+            model_name, trust_remote_code=True
+        )
+        self.model.eval()
+        self.transform_image = transforms.Compose(
+            [
+                transforms.Resize((1024, 1024)),
+                transforms.ToTensor(),
+                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+            ]
+        )
+    
+    def to(self, device: str):
+        self.model.to(device)
+
+    def cuda(self):
+        self.model.cuda()
+
+    def cpu(self):
+        self.model.cpu()
+        
+    def __call__(self, image: Image.Image) -> Image.Image:
+        image_size = image.size
+        input_images = self.transform_image(image).unsqueeze(0).to("cuda")
+        # Prediction
+        with torch.no_grad():
+            preds = self.model(input_images)[-1].sigmoid().cpu()
+        pred = preds[0].squeeze()
+        pred_pil = transforms.ToPILImage()(pred)
+        mask = pred_pil.resize(image_size)
+        image.putalpha(mask)
+        return image
+    
\ No newline at end of file
diff --git a/pixal3d/pipelines/rembg/__init__.py b/pixal3d/pipelines/rembg/__init__.py
new file mode 100644
index 0000000..fc1eed1
--- /dev/null
+++ b/pixal3d/pipelines/rembg/__init__.py
@@ -0,0 +1 @@
+from .BiRefNet import *
diff --git a/pixal3d/pipelines/samplers/__init__.py b/pixal3d/pipelines/samplers/__init__.py
new file mode 100644
index 0000000..4a69b95
--- /dev/null
+++ b/pixal3d/pipelines/samplers/__init__.py
@@ -0,0 +1,6 @@
+from .base import Sampler
+from .flow_euler import (
+    FlowEulerSampler,
+    FlowEulerCfgSampler,
+    FlowEulerGuidanceIntervalSampler,
+)
\ No newline at end of file
diff --git a/pixal3d/pipelines/samplers/base.py b/pixal3d/pipelines/samplers/base.py
new file mode 100644
index 0000000..1966ce7
--- /dev/null
+++ b/pixal3d/pipelines/samplers/base.py
@@ -0,0 +1,20 @@
+from typing import *
+from abc import ABC, abstractmethod
+
+
+class Sampler(ABC):
+    """
+    A base class for samplers.
+    """
+
+    @abstractmethod
+    def sample(
+        self,
+        model,
+        **kwargs
+    ):
+        """
+        Sample from a model.
+        """
+        pass
+    
\ No newline at end of file
diff --git a/pixal3d/pipelines/samplers/classifier_free_guidance_mixin.py b/pixal3d/pipelines/samplers/classifier_free_guidance_mixin.py
new file mode 100644
index 0000000..8c7a4da
--- /dev/null
+++ b/pixal3d/pipelines/samplers/classifier_free_guidance_mixin.py
@@ -0,0 +1,29 @@
+from typing import *
+
+
+class ClassifierFreeGuidanceSamplerMixin:
+    """
+    A mixin class for samplers that apply classifier-free guidance.
+    """
+
+    def _inference_model(self, model, x_t, t, cond, neg_cond, guidance_strength, guidance_rescale=0.0, **kwargs):
+        if guidance_strength == 1:
+            return super()._inference_model(model, x_t, t, cond, **kwargs)
+        elif guidance_strength == 0:
+            return super()._inference_model(model, x_t, t, neg_cond, **kwargs)
+        else:
+            pred_pos = super()._inference_model(model, x_t, t, cond, **kwargs)
+            pred_neg = super()._inference_model(model, x_t, t, neg_cond, **kwargs)
+            pred = guidance_strength * pred_pos + (1 - guidance_strength) * pred_neg
+            
+            # CFG rescale
+            if guidance_rescale > 0:
+                x_0_pos = self._pred_to_xstart(x_t, t, pred_pos)
+                x_0_cfg = self._pred_to_xstart(x_t, t, pred)
+                std_pos = x_0_pos.std(dim=list(range(1, x_0_pos.ndim)), keepdim=True)
+                std_cfg = x_0_cfg.std(dim=list(range(1, x_0_cfg.ndim)), keepdim=True)
+                x_0_rescaled = x_0_cfg * (std_pos / std_cfg)
+                x_0 = guidance_rescale * x_0_rescaled + (1 - guidance_rescale) * x_0_cfg
+                pred = self._xstart_to_pred(x_t, t, x_0)
+                
+            return pred
diff --git a/pixal3d/pipelines/samplers/flow_euler.py b/pixal3d/pipelines/samplers/flow_euler.py
new file mode 100644
index 0000000..5ff72b8
--- /dev/null
+++ b/pixal3d/pipelines/samplers/flow_euler.py
@@ -0,0 +1,208 @@
+from typing import *
+import torch
+import numpy as np
+from tqdm import tqdm
+from easydict import EasyDict as edict
+from .base import Sampler
+from .classifier_free_guidance_mixin import ClassifierFreeGuidanceSamplerMixin
+from .guidance_interval_mixin import GuidanceIntervalSamplerMixin
+
+
+class FlowEulerSampler(Sampler):
+    """
+    Generate samples from a flow-matching model using Euler sampling.
+
+    Args:
+        sigma_min: The minimum scale of noise in flow.
+    """
+    def __init__(
+        self,
+        sigma_min: float,
+    ):
+        self.sigma_min = sigma_min
+
+    def _eps_to_xstart(self, x_t, t, eps):
+        assert x_t.shape == eps.shape
+        return (x_t - (self.sigma_min + (1 - self.sigma_min) * t) * eps) / (1 - t)
+
+    def _xstart_to_eps(self, x_t, t, x_0):
+        assert x_t.shape == x_0.shape
+        return (x_t - (1 - t) * x_0) / (self.sigma_min + (1 - self.sigma_min) * t)
+
+    def _v_to_xstart_eps(self, x_t, t, v):
+        assert x_t.shape == v.shape
+        eps = (1 - t) * v + x_t
+        x_0 = (1 - self.sigma_min) * x_t - (self.sigma_min + (1 - self.sigma_min) * t) * v
+        return x_0, eps
+    
+    def _pred_to_xstart(self, x_t, t, pred):
+        return (1 - self.sigma_min) * x_t - (self.sigma_min + (1 - self.sigma_min) * t) * pred
+
+    def _xstart_to_pred(self, x_t, t, x_0):
+        return ((1 - self.sigma_min) * x_t - x_0) / (self.sigma_min + (1 - self.sigma_min) * t)
+
+    def _inference_model(self, model, x_t, t, cond=None, **kwargs):
+        t = torch.tensor([1000 * t] * x_t.shape[0], device=x_t.device, dtype=torch.float32)
+        return model(x_t, t, cond, **kwargs)
+
+    def _get_model_prediction(self, model, x_t, t, cond=None, **kwargs):
+        pred_v = self._inference_model(model, x_t, t, cond, **kwargs)
+        pred_x_0, pred_eps = self._v_to_xstart_eps(x_t=x_t, t=t, v=pred_v)
+        return pred_x_0, pred_eps, pred_v
+
+    @torch.no_grad()
+    def sample_once(
+        self,
+        model,
+        x_t,
+        t: float,
+        t_prev: float,
+        cond: Optional[Any] = None,
+        **kwargs
+    ):
+        """
+        Sample x_{t-1} from the model using Euler method.
+        
+        Args:
+            model: The model to sample from.
+            x_t: The [N x C x ...] tensor of noisy inputs at time t.
+            t: The current timestep.
+            t_prev: The previous timestep.
+            cond: conditional information.
+            **kwargs: Additional arguments for model inference.
+
+        Returns:
+            a dict containing the following
+            - 'pred_x_prev': x_{t-1}.
+            - 'pred_x_0': a prediction of x_0.
+        """
+        pred_x_0, pred_eps, pred_v = self._get_model_prediction(model, x_t, t, cond, **kwargs)
+        pred_x_prev = x_t - (t - t_prev) * pred_v
+        return edict({"pred_x_prev": pred_x_prev, "pred_x_0": pred_x_0})
+
+    @torch.no_grad()
+    def sample(
+        self,
+        model,
+        noise,
+        cond: Optional[Any] = None,
+        steps: int = 50,
+        rescale_t: float = 1.0,
+        verbose: bool = True,
+        tqdm_desc: str = "Sampling",
+        **kwargs
+    ):
+        """
+        Generate samples from the model using Euler method.
+        
+        Args:
+            model: The model to sample from.
+            noise: The initial noise tensor.
+            cond: conditional information.
+            steps: The number of steps to sample.
+            rescale_t: The rescale factor for t.
+            verbose: If True, show a progress bar.
+            tqdm_desc: A customized tqdm desc.
+            **kwargs: Additional arguments for model_inference.
+
+        Returns:
+            a dict containing the following
+            - 'samples': the model samples.
+            - 'pred_x_t': a list of prediction of x_t.
+            - 'pred_x_0': a list of prediction of x_0.
+        """
+        sample = noise
+        t_seq = np.linspace(1, 0, steps + 1)
+        t_seq = rescale_t * t_seq / (1 + (rescale_t - 1) * t_seq)
+        t_seq = t_seq.tolist()
+        t_pairs = list((t_seq[i], t_seq[i + 1]) for i in range(steps))
+        ret = edict({"samples": None, "pred_x_t": [], "pred_x_0": []})
+        for t, t_prev in tqdm(t_pairs, desc=tqdm_desc, disable=not verbose):
+            out = self.sample_once(model, sample, t, t_prev, cond, **kwargs)
+            sample = out.pred_x_prev
+            ret.pred_x_t.append(out.pred_x_prev)
+            ret.pred_x_0.append(out.pred_x_0)
+        ret.samples = sample
+        return ret
+
+
+class FlowEulerCfgSampler(ClassifierFreeGuidanceSamplerMixin, FlowEulerSampler):
+    """
+    Generate samples from a flow-matching model using Euler sampling with classifier-free guidance.
+    """
+    @torch.no_grad()
+    def sample(
+        self,
+        model,
+        noise,
+        cond,
+        neg_cond,
+        steps: int = 50,
+        rescale_t: float = 1.0,
+        guidance_strength: float = 3.0,
+        verbose: bool = True,
+        **kwargs
+    ):
+        """
+        Generate samples from the model using Euler method.
+        
+        Args:
+            model: The model to sample from.
+            noise: The initial noise tensor.
+            cond: conditional information.
+            neg_cond: negative conditional information.
+            steps: The number of steps to sample.
+            rescale_t: The rescale factor for t.
+            guidance_strength: The strength of classifier-free guidance.
+            verbose: If True, show a progress bar.
+            **kwargs: Additional arguments for model_inference.
+
+        Returns:
+            a dict containing the following
+            - 'samples': the model samples.
+            - 'pred_x_t': a list of prediction of x_t.
+            - 'pred_x_0': a list of prediction of x_0.
+        """
+        return super().sample(model, noise, cond, steps, rescale_t, verbose, neg_cond=neg_cond, guidance_strength=guidance_strength, **kwargs)
+
+
+class FlowEulerGuidanceIntervalSampler(GuidanceIntervalSamplerMixin, ClassifierFreeGuidanceSamplerMixin, FlowEulerSampler):
+    """
+    Generate samples from a flow-matching model using Euler sampling with classifier-free guidance and interval.
+    """
+    @torch.no_grad()
+    def sample(
+        self,
+        model,
+        noise,
+        cond,
+        neg_cond,
+        steps: int = 50,
+        rescale_t: float = 1.0,
+        guidance_strength: float = 3.0,
+        guidance_interval: Tuple[float, float] = (0.0, 1.0),
+        verbose: bool = True,
+        **kwargs
+    ):
+        """
+        Generate samples from the model using Euler method.
+        
+        Args:
+            model: The model to sample from.
+            noise: The initial noise tensor.
+            cond: conditional information.
+            neg_cond: negative conditional information.
+            steps: The number of steps to sample.
+            rescale_t: The rescale factor for t.
+            guidance_strength: The strength of classifier-free guidance.
+            guidance_interval: The interval for classifier-free guidance.
+            verbose: If True, show a progress bar.
+            **kwargs: Additional arguments for model_inference.
+
+        Returns:
+            a dict containing the following
+            - 'samples': the model samples.
+            - 'pred_x_t': a list of prediction of x_t.
+            - 'pred_x_0': a list of prediction of x_0.
+        """
+        return super().sample(model, noise, cond, steps, rescale_t, verbose, neg_cond=neg_cond, guidance_strength=guidance_strength, guidance_interval=guidance_interval, **kwargs)
diff --git a/pixal3d/pipelines/samplers/guidance_interval_mixin.py b/pixal3d/pipelines/samplers/guidance_interval_mixin.py
new file mode 100644
index 0000000..3f57869
--- /dev/null
+++ b/pixal3d/pipelines/samplers/guidance_interval_mixin.py
@@ -0,0 +1,13 @@
+from typing import *
+
+
+class GuidanceIntervalSamplerMixin:
+    """
+    A mixin class for samplers that apply classifier-free guidance with interval.
+    """
+
+    def _inference_model(self, model, x_t, t, cond, guidance_strength, guidance_interval, **kwargs):
+        if guidance_interval[0] <= t <= guidance_interval[1]:
+            return super()._inference_model(model, x_t, t, cond, guidance_strength=guidance_strength, **kwargs)
+        else:
+            return super()._inference_model(model, x_t, t, cond, guidance_strength=1, **kwargs)
diff --git a/pixal3d/pipelines/trellis2_image_to_3d.py b/pixal3d/pipelines/trellis2_image_to_3d.py
new file mode 100644
index 0000000..e4bffc2
--- /dev/null
+++ b/pixal3d/pipelines/trellis2_image_to_3d.py
@@ -0,0 +1,610 @@
+from typing import *
+import torch
+import torch.nn as nn
+import numpy as np
+from PIL import Image
+from .base import Pipeline
+from . import samplers, rembg
+from ..modules.sparse import SparseTensor
+from ..modules import image_feature_extractor
+from ..representations import Mesh, MeshWithVoxel
+
+
+class Trellis2ImageTo3DPipeline(Pipeline):
+    """
+    Pipeline for inferring Trellis2 image-to-3D models.
+
+    Args:
+        models (dict[str, nn.Module]): The models to use in the pipeline.
+        sparse_structure_sampler (samplers.Sampler): The sampler for the sparse structure.
+        shape_slat_sampler (samplers.Sampler): The sampler for the structured latent.
+        tex_slat_sampler (samplers.Sampler): The sampler for the texture latent.
+        sparse_structure_sampler_params (dict): The parameters for the sparse structure sampler.
+        shape_slat_sampler_params (dict): The parameters for the structured latent sampler.
+        tex_slat_sampler_params (dict): The parameters for the texture latent sampler.
+        shape_slat_normalization (dict): The normalization parameters for the structured latent.
+        tex_slat_normalization (dict): The normalization parameters for the texture latent.
+        image_cond_model (Callable): The image conditioning model.
+        rembg_model (Callable): The model for removing background.
+        low_vram (bool): Whether to use low-VRAM mode.
+    """
+    model_names_to_load = [
+        'sparse_structure_flow_model',
+        'sparse_structure_decoder',
+        'shape_slat_flow_model_512',
+        'shape_slat_flow_model_1024',
+        'shape_slat_decoder',
+        'tex_slat_flow_model_512',
+        'tex_slat_flow_model_1024',
+        'tex_slat_decoder',
+    ]
+
+    def __init__(
+        self,
+        models: dict[str, nn.Module] = None,
+        sparse_structure_sampler: samplers.Sampler = None,
+        shape_slat_sampler: samplers.Sampler = None,
+        tex_slat_sampler: samplers.Sampler = None,
+        sparse_structure_sampler_params: dict = None,
+        shape_slat_sampler_params: dict = None,
+        tex_slat_sampler_params: dict = None,
+        shape_slat_normalization: dict = None,
+        tex_slat_normalization: dict = None,
+        image_cond_model: Callable = None,
+        rembg_model: Callable = None,
+        low_vram: bool = True,
+        default_pipeline_type: str = '1024_cascade',
+    ):
+        if models is None:
+            return
+        super().__init__(models)
+        self.sparse_structure_sampler = sparse_structure_sampler
+        self.shape_slat_sampler = shape_slat_sampler
+        self.tex_slat_sampler = tex_slat_sampler
+        self.sparse_structure_sampler_params = sparse_structure_sampler_params
+        self.shape_slat_sampler_params = shape_slat_sampler_params
+        self.tex_slat_sampler_params = tex_slat_sampler_params
+        self.shape_slat_normalization = shape_slat_normalization
+        self.tex_slat_normalization = tex_slat_normalization
+        self.image_cond_model = image_cond_model
+        self.rembg_model = rembg_model
+        self.low_vram = low_vram
+        self.default_pipeline_type = default_pipeline_type
+        self.pbr_attr_layout = {
+            'base_color': slice(0, 3),
+            'metallic': slice(3, 4),
+            'roughness': slice(4, 5),
+            'alpha': slice(5, 6),
+        }
+        self._device = 'cpu'
+
+    @classmethod
+    def from_pretrained(cls, path: str, config_file: str = "pipeline.json") -> "Trellis2ImageTo3DPipeline":
+        """
+        Load a pretrained model.
+
+        Args:
+            path (str): The path to the model. Can be either local path or a Hugging Face repository.
+        """
+        pipeline = super().from_pretrained(path, config_file)
+        args = pipeline._pretrained_args
+
+        pipeline.sparse_structure_sampler = getattr(samplers, args['sparse_structure_sampler']['name'])(**args['sparse_structure_sampler']['args'])
+        pipeline.sparse_structure_sampler_params = args['sparse_structure_sampler']['params']
+
+        pipeline.shape_slat_sampler = getattr(samplers, args['shape_slat_sampler']['name'])(**args['shape_slat_sampler']['args'])
+        pipeline.shape_slat_sampler_params = args['shape_slat_sampler']['params']
+
+        pipeline.tex_slat_sampler = getattr(samplers, args['tex_slat_sampler']['name'])(**args['tex_slat_sampler']['args'])
+        pipeline.tex_slat_sampler_params = args['tex_slat_sampler']['params']
+
+        pipeline.shape_slat_normalization = args['shape_slat_normalization']
+        pipeline.tex_slat_normalization = args['tex_slat_normalization']
+
+        # HACK: Replace dinov3 model source with camenduru mirror
+        image_cond_args = args['image_cond_model']['args'].copy()
+        if image_cond_args.get('model_name') == 'facebook/dinov3-vitl16-pretrain-lvd1689m':
+            image_cond_args['model_name'] = 'camenduru/dinov3-vitl16-pretrain-lvd1689m'
+        pipeline.image_cond_model = getattr(image_feature_extractor, args['image_cond_model']['name'])(**image_cond_args)
+        pipeline.rembg_model = getattr(rembg, args['rembg_model']['name'])(**args['rembg_model']['args'])
+        
+        pipeline.low_vram = args.get('low_vram', True)
+        pipeline.default_pipeline_type = args.get('default_pipeline_type', '1024_cascade')
+        pipeline.pbr_attr_layout = {
+            'base_color': slice(0, 3),
+            'metallic': slice(3, 4),
+            'roughness': slice(4, 5),
+            'alpha': slice(5, 6),
+        }
+        pipeline._device = 'cpu'
+
+        return pipeline
+
+    def to(self, device: torch.device) -> None:
+        self._device = device
+        if not self.low_vram:
+            super().to(device)
+            self.image_cond_model.to(device)
+            if self.rembg_model is not None:
+                self.rembg_model.to(device)
+
+    def preprocess_image(self, input: Image.Image, bg_color: tuple = (0, 0, 0)) -> Image.Image:
+        """
+        Preprocess the input image.
+
+        Args:
+            input: Input image (RGB or RGBA).
+            bg_color: Background color (R, G, B) in 0~255. Default black (0,0,0).
+        """
+        # if has alpha channel, use it directly; otherwise, remove background
+        has_alpha = False
+        if input.mode == 'RGBA':
+            alpha = np.array(input)[:, :, 3]
+            if not np.all(alpha == 255):
+                has_alpha = True
+        max_size = max(input.size)
+        scale = min(1, 1024 / max_size)
+        if scale < 1:
+            input = input.resize((int(input.width * scale), int(input.height * scale)), Image.Resampling.LANCZOS)
+        if has_alpha:
+            output = input
+        else:
+            input = input.convert('RGB')
+            if self.low_vram:
+                self.rembg_model.to(self.device)
+            output = self.rembg_model(input)
+            if self.low_vram:
+                self.rembg_model.cpu()
+        output_np = np.array(output)
+        alpha = output_np[:, :, 3]
+        bbox = np.argwhere(alpha > 0.8 * 255)
+        bbox = np.min(bbox[:, 1]), np.min(bbox[:, 0]), np.max(bbox[:, 1]), np.max(bbox[:, 0])
+        center = (bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2
+        size = max(bbox[2] - bbox[0], bbox[3] - bbox[1])
+        size = int(size * 1.1)
+        bbox = center[0] - size // 2, center[1] - size // 2, center[0] + size // 2, center[1] + size // 2
+        output = output.crop(bbox)  # type: ignore
+        output = np.array(output).astype(np.float32) / 255
+        rgb = output[:, :, :3]
+        a = output[:, :, 3:4]
+        bg = np.array(bg_color, dtype=np.float32) / 255.0
+        output = rgb * a + bg * (1.0 - a)
+        output = Image.fromarray((np.clip(output, 0, 1) * 255).astype(np.uint8))
+        return output
+        
+    def get_cond(self, image: Union[torch.Tensor, list[Image.Image]], resolution: int, include_neg_cond: bool = True) -> dict:
+        """
+        Get the conditioning information for the model.
+
+        Args:
+            image (Union[torch.Tensor, list[Image.Image]]): The image prompts.
+
+        Returns:
+            dict: The conditioning information
+        """
+        self.image_cond_model.image_size = resolution
+        if self.low_vram:
+            self.image_cond_model.to(self.device)
+        cond = self.image_cond_model(image)
+        if self.low_vram:
+            self.image_cond_model.cpu()
+        if not include_neg_cond:
+            return {'cond': cond}
+        neg_cond = torch.zeros_like(cond)
+        return {
+            'cond': cond,
+            'neg_cond': neg_cond,
+        }
+
+    def sample_sparse_structure(
+        self,
+        cond: dict,
+        resolution: int,
+        num_samples: int = 1,
+        sampler_params: dict = {},
+    ) -> torch.Tensor:
+        """
+        Sample sparse structures with the given conditioning.
+        
+        Args:
+            cond (dict): The conditioning information.
+            resolution (int): The resolution of the sparse structure.
+            num_samples (int): The number of samples to generate.
+            sampler_params (dict): Additional parameters for the sampler.
+        """
+        # Sample sparse structure latent
+        flow_model = self.models['sparse_structure_flow_model']
+        reso = flow_model.resolution
+        in_channels = flow_model.in_channels
+        noise = torch.randn(num_samples, in_channels, reso, reso, reso).to(self.device)
+        sampler_params = {**self.sparse_structure_sampler_params, **sampler_params}
+        if self.low_vram:
+            flow_model.to(self.device)
+        z_s = self.sparse_structure_sampler.sample(
+            flow_model,
+            noise,
+            **cond,
+            **sampler_params,
+            verbose=True,
+            tqdm_desc="Sampling sparse structure",
+        ).samples
+        if self.low_vram:
+            flow_model.cpu()
+        
+        # Decode sparse structure latent
+        decoder = self.models['sparse_structure_decoder']
+        if self.low_vram:
+            decoder.to(self.device)
+        decoded = decoder(z_s)>0
+        if self.low_vram:
+            decoder.cpu()
+        if resolution != decoded.shape[2]:
+            ratio = decoded.shape[2] // resolution
+            decoded = torch.nn.functional.max_pool3d(decoded.float(), ratio, ratio, 0) > 0.5
+        coords = torch.argwhere(decoded)[:, [0, 2, 3, 4]].int()
+
+        return coords
+
+    def sample_shape_slat(
+        self,
+        cond: dict,
+        flow_model,
+        coords: torch.Tensor,
+        sampler_params: dict = {},
+    ) -> SparseTensor:
+        """
+        Sample structured latent with the given conditioning.
+        
+        Args:
+            cond (dict): The conditioning information.
+            coords (torch.Tensor): The coordinates of the sparse structure.
+            sampler_params (dict): Additional parameters for the sampler.
+        """
+        # Sample structured latent
+        noise = SparseTensor(
+            feats=torch.randn(coords.shape[0], flow_model.in_channels).to(self.device),
+            coords=coords,
+        )
+        sampler_params = {**self.shape_slat_sampler_params, **sampler_params}
+        if self.low_vram:
+            flow_model.to(self.device)
+        slat = self.shape_slat_sampler.sample(
+            flow_model,
+            noise,
+            **cond,
+            **sampler_params,
+            verbose=True,
+            tqdm_desc="Sampling shape SLat",
+        ).samples
+        if self.low_vram:
+            flow_model.cpu()
+
+        std = torch.tensor(self.shape_slat_normalization['std'])[None].to(slat.device)
+        mean = torch.tensor(self.shape_slat_normalization['mean'])[None].to(slat.device)
+        slat = slat * std + mean
+        
+        return slat
+    
+    def sample_shape_slat_cascade(
+        self,
+        lr_cond: dict,
+        cond: dict,
+        flow_model_lr,
+        flow_model,
+        lr_resolution: int,
+        resolution: int,
+        coords: torch.Tensor,
+        sampler_params: dict = {},
+        max_num_tokens: int = 49152,
+    ) -> SparseTensor:
+        """
+        Sample structured latent with the given conditioning.
+        
+        Args:
+            cond (dict): The conditioning information.
+            coords (torch.Tensor): The coordinates of the sparse structure.
+            sampler_params (dict): Additional parameters for the sampler.
+        """
+        # LR
+        noise = SparseTensor(
+            feats=torch.randn(coords.shape[0], flow_model_lr.in_channels).to(self.device),
+            coords=coords,
+        )
+        sampler_params = {**self.shape_slat_sampler_params, **sampler_params}
+        if self.low_vram:
+            flow_model_lr.to(self.device)
+        slat = self.shape_slat_sampler.sample(
+            flow_model_lr,
+            noise,
+            **lr_cond,
+            **sampler_params,
+            verbose=True,
+            tqdm_desc="Sampling shape SLat",
+        ).samples
+        if self.low_vram:
+            flow_model_lr.cpu()
+        std = torch.tensor(self.shape_slat_normalization['std'])[None].to(slat.device)
+        mean = torch.tensor(self.shape_slat_normalization['mean'])[None].to(slat.device)
+        slat = slat * std + mean
+        
+        # Upsample
+        if self.low_vram:
+            self.models['shape_slat_decoder'].to(self.device)
+            self.models['shape_slat_decoder'].low_vram = True
+        hr_coords = self.models['shape_slat_decoder'].upsample(slat, upsample_times=4)
+        if self.low_vram:
+            self.models['shape_slat_decoder'].cpu()
+            self.models['shape_slat_decoder'].low_vram = False
+        hr_resolution = resolution
+        while True:
+            quant_coords = torch.cat([
+                hr_coords[:, :1],
+                ((hr_coords[:, 1:] + 0.5) / lr_resolution * (hr_resolution // 16)).int(),
+            ], dim=1)
+            coords = quant_coords.unique(dim=0)
+            num_tokens = coords.shape[0]
+            if num_tokens < max_num_tokens or hr_resolution == 1024:
+                if hr_resolution != resolution:
+                    print(f"Due to the limited number of tokens, the resolution is reduced to {hr_resolution}.")
+                break
+            hr_resolution -= 128
+        
+        # Sample structured latent
+        noise = SparseTensor(
+            feats=torch.randn(coords.shape[0], flow_model.in_channels).to(self.device),
+            coords=coords,
+        )
+        sampler_params = {**self.shape_slat_sampler_params, **sampler_params}
+        if self.low_vram:
+            flow_model.to(self.device)
+        slat = self.shape_slat_sampler.sample(
+            flow_model,
+            noise,
+            **cond,
+            **sampler_params,
+            verbose=True,
+            tqdm_desc="Sampling shape SLat",
+        ).samples
+        if self.low_vram:
+            flow_model.cpu()
+
+        std = torch.tensor(self.shape_slat_normalization['std'])[None].to(slat.device)
+        mean = torch.tensor(self.shape_slat_normalization['mean'])[None].to(slat.device)
+        slat = slat * std + mean
+        
+        return slat, hr_resolution
+
+    def decode_shape_slat(
+        self,
+        slat: SparseTensor,
+        resolution: int,
+    ) -> Tuple[List[Mesh], List[SparseTensor]]:
+        """
+        Decode the structured latent.
+
+        Args:
+            slat (SparseTensor): The structured latent.
+
+        Returns:
+            List[Mesh]: The decoded meshes.
+            List[SparseTensor]: The decoded substructures.
+        """
+        self.models['shape_slat_decoder'].set_resolution(resolution)
+        if self.low_vram:
+            self.models['shape_slat_decoder'].to(self.device)
+            self.models['shape_slat_decoder'].low_vram = True
+        ret = self.models['shape_slat_decoder'](slat, return_subs=True)
+        if self.low_vram:
+            self.models['shape_slat_decoder'].cpu()
+            self.models['shape_slat_decoder'].low_vram = False
+        return ret
+    
+    def sample_tex_slat(
+        self,
+        cond: dict,
+        flow_model,
+        shape_slat: SparseTensor,
+        sampler_params: dict = {},
+    ) -> SparseTensor:
+        """
+        Sample structured latent with the given conditioning.
+        
+        Args:
+            cond (dict): The conditioning information.
+            shape_slat (SparseTensor): The structured latent for shape
+            sampler_params (dict): Additional parameters for the sampler.
+        """
+        # Sample structured latent
+        std = torch.tensor(self.shape_slat_normalization['std'])[None].to(shape_slat.device)
+        mean = torch.tensor(self.shape_slat_normalization['mean'])[None].to(shape_slat.device)
+        shape_slat = (shape_slat - mean) / std
+
+        in_channels = flow_model.in_channels if isinstance(flow_model, nn.Module) else flow_model[0].in_channels
+        noise = shape_slat.replace(feats=torch.randn(shape_slat.coords.shape[0], in_channels - shape_slat.feats.shape[1]).to(self.device))
+        sampler_params = {**self.tex_slat_sampler_params, **sampler_params}
+        if self.low_vram:
+            flow_model.to(self.device)
+        slat = self.tex_slat_sampler.sample(
+            flow_model,
+            noise,
+            concat_cond=shape_slat,
+            **cond,
+            **sampler_params,
+            verbose=True,
+            tqdm_desc="Sampling texture SLat",
+        ).samples
+        if self.low_vram:
+            flow_model.cpu()
+
+        std = torch.tensor(self.tex_slat_normalization['std'])[None].to(slat.device)
+        mean = torch.tensor(self.tex_slat_normalization['mean'])[None].to(slat.device)
+        slat = slat * std + mean
+        
+        return slat
+
+    def decode_tex_slat(
+        self,
+        slat: SparseTensor,
+        subs: List[SparseTensor],
+    ) -> SparseTensor:
+        """
+        Decode the structured latent.
+
+        Args:
+            slat (SparseTensor): The structured latent.
+
+        Returns:
+            SparseTensor: The decoded texture voxels
+        """
+        if self.low_vram:
+            self.models['tex_slat_decoder'].to(self.device)
+        ret = self.models['tex_slat_decoder'](slat, guide_subs=subs) * 0.5 + 0.5
+        if self.low_vram:
+            self.models['tex_slat_decoder'].cpu()
+        return ret
+    
+    @torch.no_grad()
+    def decode_latent(
+        self,
+        shape_slat: SparseTensor,
+        tex_slat: SparseTensor,
+        resolution: int,
+    ) -> List[MeshWithVoxel]:
+        """
+        Decode the latent codes.
+
+        Args:
+            shape_slat (SparseTensor): The structured latent for shape.
+            tex_slat (SparseTensor): The structured latent for texture.
+            resolution (int): The resolution of the output.
+        """
+        meshes, subs = self.decode_shape_slat(shape_slat, resolution)
+        tex_voxels = self.decode_tex_slat(tex_slat, subs)
+        out_mesh = []
+        torch.cuda.synchronize()
+        for m, v in zip(meshes, tex_voxels):
+            # try:
+            m.fill_holes()
+            # except RuntimeError as e:
+                # print(f"[WARNING] fill_holes failed (likely PyTorch/cumesh compatibility issue), skipping: {e}")
+            out_mesh.append(
+                MeshWithVoxel(
+                    m.vertices, m.faces,
+                    origin = [-0.5, -0.5, -0.5],
+                    voxel_size = 1 / resolution,
+                    coords = v.coords[:, 1:],
+                    attrs = v.feats,
+                    voxel_shape = torch.Size([*v.shape, *v.spatial_shape]),
+                    layout=self.pbr_attr_layout
+                )
+            )
+        return out_mesh
+    
+    @torch.no_grad()
+    def run(
+        self,
+        image: Image.Image,
+        num_samples: int = 1,
+        seed: int = 42,
+        sparse_structure_sampler_params: dict = {},
+        shape_slat_sampler_params: dict = {},
+        tex_slat_sampler_params: dict = {},
+        preprocess_image: bool = True,
+        return_latent: bool = False,
+        pipeline_type: Optional[str] = None,
+        max_num_tokens: int = 49152,
+    ) -> List[MeshWithVoxel]:
+        """
+        Run the pipeline.
+
+        Args:
+            image (Image.Image): The image prompt.
+            num_samples (int): The number of samples to generate.
+            seed (int): The random seed.
+            sparse_structure_sampler_params (dict): Additional parameters for the sparse structure sampler.
+            shape_slat_sampler_params (dict): Additional parameters for the shape SLat sampler.
+            tex_slat_sampler_params (dict): Additional parameters for the texture SLat sampler.
+            preprocess_image (bool): Whether to preprocess the image.
+            return_latent (bool): Whether to return the latent codes.
+            pipeline_type (str): The type of the pipeline. Options: '512', '1024', '1024_cascade', '1536_cascade'.
+            max_num_tokens (int): The maximum number of tokens to use.
+        """
+        # Check pipeline type
+        pipeline_type = pipeline_type or self.default_pipeline_type
+        if pipeline_type == '512':
+            assert 'shape_slat_flow_model_512' in self.models, "No 512 resolution shape SLat flow model found."
+            assert 'tex_slat_flow_model_512' in self.models, "No 512 resolution texture SLat flow model found."
+        elif pipeline_type == '1024':
+            assert 'shape_slat_flow_model_1024' in self.models, "No 1024 resolution shape SLat flow model found."
+            assert 'tex_slat_flow_model_1024' in self.models, "No 1024 resolution texture SLat flow model found."
+        elif pipeline_type == '1024_cascade':
+            assert 'shape_slat_flow_model_512' in self.models, "No 512 resolution shape SLat flow model found."
+            assert 'shape_slat_flow_model_1024' in self.models, "No 1024 resolution shape SLat flow model found."
+            assert 'tex_slat_flow_model_1024' in self.models, "No 1024 resolution texture SLat flow model found."
+        elif pipeline_type == '1536_cascade':
+            assert 'shape_slat_flow_model_512' in self.models, "No 512 resolution shape SLat flow model found."
+            assert 'shape_slat_flow_model_1024' in self.models, "No 1024 resolution shape SLat flow model found."
+            assert 'tex_slat_flow_model_1024' in self.models, "No 1024 resolution texture SLat flow model found."
+        else:
+            raise ValueError(f"Invalid pipeline type: {pipeline_type}")
+        
+        if preprocess_image:
+            image = self.preprocess_image(image)
+        torch.manual_seed(seed)
+        cond_512 = self.get_cond([image], 512)
+        cond_1024 = self.get_cond([image], 1024) if pipeline_type != '512' else None
+        ss_res = {'512': 32, '1024': 64, '1024_cascade': 32, '1536_cascade': 32}[pipeline_type]
+        coords = self.sample_sparse_structure(
+            cond_512, ss_res,
+            num_samples, sparse_structure_sampler_params
+        )
+        if pipeline_type == '512':
+            shape_slat = self.sample_shape_slat(
+                cond_512, self.models['shape_slat_flow_model_512'],
+                coords, shape_slat_sampler_params
+            )
+            tex_slat = self.sample_tex_slat(
+                cond_512, self.models['tex_slat_flow_model_512'],
+                shape_slat, tex_slat_sampler_params
+            )
+            res = 512
+        elif pipeline_type == '1024':
+            shape_slat = self.sample_shape_slat(
+                cond_1024, self.models['shape_slat_flow_model_1024'],
+                coords, shape_slat_sampler_params
+            )
+            tex_slat = self.sample_tex_slat(
+                cond_1024, self.models['tex_slat_flow_model_1024'],
+                shape_slat, tex_slat_sampler_params
+            )
+            res = 1024
+        elif pipeline_type == '1024_cascade':
+            shape_slat, res = self.sample_shape_slat_cascade(
+                cond_512, cond_1024,
+                self.models['shape_slat_flow_model_512'], self.models['shape_slat_flow_model_1024'],
+                512, 1024,
+                coords, shape_slat_sampler_params,
+                max_num_tokens
+            )
+            tex_slat = self.sample_tex_slat(
+                cond_1024, self.models['tex_slat_flow_model_1024'],
+                shape_slat, tex_slat_sampler_params
+            )
+        elif pipeline_type == '1536_cascade':
+            shape_slat, res = self.sample_shape_slat_cascade(
+                cond_512, cond_1024,
+                self.models['shape_slat_flow_model_512'], self.models['shape_slat_flow_model_1024'],
+                512, 1536,
+                coords, shape_slat_sampler_params,
+                max_num_tokens
+            )
+            tex_slat = self.sample_tex_slat(
+                cond_1024, self.models['tex_slat_flow_model_1024'],
+                shape_slat, tex_slat_sampler_params
+            )
+        torch.cuda.empty_cache()
+        out_mesh = self.decode_latent(shape_slat, tex_slat, res)
+        if return_latent:
+            return out_mesh, (shape_slat, tex_slat, res)
+        else:
+            return out_mesh
diff --git a/pixal3d/pipelines/trellis2_texturing.py b/pixal3d/pipelines/trellis2_texturing.py
new file mode 100644
index 0000000..c184b5e
--- /dev/null
+++ b/pixal3d/pipelines/trellis2_texturing.py
@@ -0,0 +1,408 @@
+from typing import *
+import torch
+import torch.nn as nn
+import numpy as np
+from PIL import Image
+import trimesh
+from .base import Pipeline
+from . import samplers, rembg
+from ..modules.sparse import SparseTensor
+from ..modules import image_feature_extractor
+import o_voxel
+import cumesh
+import nvdiffrast.torch as dr
+import cv2
+import flex_gemm
+
+
+class Trellis2TexturingPipeline(Pipeline):
+    """
+    Pipeline for inferring Trellis2 image-to-3D models.
+
+    Args:
+        models (dict[str, nn.Module]): The models to use in the pipeline.
+        tex_slat_sampler (samplers.Sampler): The sampler for the texture latent.
+        tex_slat_sampler_params (dict): The parameters for the texture latent sampler.
+        shape_slat_normalization (dict): The normalization parameters for the structured latent.
+        tex_slat_normalization (dict): The normalization parameters for the texture latent.
+        image_cond_model (Callable): The image conditioning model.
+        rembg_model (Callable): The model for removing background.
+        low_vram (bool): Whether to use low-VRAM mode.
+    """
+    model_names_to_load = [
+        'shape_slat_encoder',
+        'tex_slat_decoder',
+        'tex_slat_flow_model_512',
+        'tex_slat_flow_model_1024'
+    ]
+
+    def __init__(
+        self,
+        models: dict[str, nn.Module] = None,
+        tex_slat_sampler: samplers.Sampler = None,
+        tex_slat_sampler_params: dict = None,
+        shape_slat_normalization: dict = None,
+        tex_slat_normalization: dict = None,
+        image_cond_model: Callable = None,
+        rembg_model: Callable = None,
+        low_vram: bool = True,
+    ):
+        if models is None:
+            return
+        super().__init__(models)
+        self.tex_slat_sampler = tex_slat_sampler
+        self.tex_slat_sampler_params = tex_slat_sampler_params
+        self.shape_slat_normalization = shape_slat_normalization
+        self.tex_slat_normalization = tex_slat_normalization
+        self.image_cond_model = image_cond_model
+        self.rembg_model = rembg_model
+        self.low_vram = low_vram
+        self.pbr_attr_layout = {
+            'base_color': slice(0, 3),
+            'metallic': slice(3, 4),
+            'roughness': slice(4, 5),
+            'alpha': slice(5, 6),
+        }
+        self._device = 'cpu'
+
+    @classmethod
+    def from_pretrained(cls, path: str, config_file: str = "pipeline.json") -> "Trellis2TexturingPipeline":
+        """
+        Load a pretrained model.
+
+        Args:
+            path (str): The path to the model. Can be either local path or a Hugging Face repository.
+        """
+        pipeline = super().from_pretrained(path, config_file)
+        args = pipeline._pretrained_args
+
+        pipeline.tex_slat_sampler = getattr(samplers, args['tex_slat_sampler']['name'])(**args['tex_slat_sampler']['args'])
+        pipeline.tex_slat_sampler_params = args['tex_slat_sampler']['params']
+
+        pipeline.shape_slat_normalization = args['shape_slat_normalization']
+        pipeline.tex_slat_normalization = args['tex_slat_normalization']
+
+        pipeline.image_cond_model = getattr(image_feature_extractor, args['image_cond_model']['name'])(**args['image_cond_model']['args'])
+        pipeline.rembg_model = getattr(rembg, args['rembg_model']['name'])(**args['rembg_model']['args'])
+
+        pipeline.low_vram = args.get('low_vram', True)
+        pipeline.pbr_attr_layout = {
+            'base_color': slice(0, 3),
+            'metallic': slice(3, 4),
+            'roughness': slice(4, 5),
+            'alpha': slice(5, 6),
+        }
+        pipeline._device = 'cpu'
+        return pipeline
+
+    def to(self, device: torch.device) -> None:
+        self._device = device
+        if not self.low_vram:
+            super().to(device)
+            self.image_cond_model.to(device)
+            if self.rembg_model is not None:
+                self.rembg_model.to(device)
+
+    def preprocess_mesh(self, mesh: trimesh.Trimesh) -> trimesh.Trimesh:
+        """
+        Preprocess the input mesh.
+        """
+        vertices = mesh.vertices
+        vertices_min = vertices.min(axis=0)
+        vertices_max = vertices.max(axis=0)
+        center = (vertices_min + vertices_max) / 2
+        scale = 0.99999 / (vertices_max - vertices_min).max()
+        vertices = (vertices - center) * scale
+        tmp = vertices[:, 1].copy()
+        vertices[:, 1] = -vertices[:, 2]
+        vertices[:, 2] = tmp
+        assert np.all(vertices >= -0.5) and np.all(vertices <= 0.5), 'vertices out of range'
+        return trimesh.Trimesh(vertices=vertices, faces=mesh.faces, process=False)
+
+    def preprocess_image(self, input: Image.Image) -> Image.Image:
+        """
+        Preprocess the input image.
+        """
+        # if has alpha channel, use it directly; otherwise, remove background
+        has_alpha = False
+        if input.mode == 'RGBA':
+            alpha = np.array(input)[:, :, 3]
+            if not np.all(alpha == 255):
+                has_alpha = True
+        max_size = max(input.size)
+        scale = min(1, 1024 / max_size)
+        if scale < 1:
+            input = input.resize((int(input.width * scale), int(input.height * scale)), Image.Resampling.LANCZOS)
+        if has_alpha:
+            output = input
+        else:
+            input = input.convert('RGB')
+            if self.low_vram:
+                self.rembg_model.to(self.device)
+            output = self.rembg_model(input)
+            if self.low_vram:
+                self.rembg_model.cpu()
+        output_np = np.array(output)
+        alpha = output_np[:, :, 3]
+        bbox = np.argwhere(alpha > 0.8 * 255)
+        bbox = np.min(bbox[:, 1]), np.min(bbox[:, 0]), np.max(bbox[:, 1]), np.max(bbox[:, 0])
+        center = (bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2
+        size = max(bbox[2] - bbox[0], bbox[3] - bbox[1])
+        size = int(size * 1)
+        bbox = center[0] - size // 2, center[1] - size // 2, center[0] + size // 2, center[1] + size // 2
+        output = output.crop(bbox)  # type: ignore
+        output = np.array(output).astype(np.float32) / 255
+        output = output[:, :, :3] * output[:, :, 3:4]
+        output = Image.fromarray((output * 255).astype(np.uint8))
+        return output
+        
+    def get_cond(self, image: Union[torch.Tensor, list[Image.Image]], resolution: int, include_neg_cond: bool = True) -> dict:
+        """
+        Get the conditioning information for the model.
+
+        Args:
+            image (Union[torch.Tensor, list[Image.Image]]): The image prompts.
+
+        Returns:
+            dict: The conditioning information
+        """
+        self.image_cond_model.image_size = resolution
+        if self.low_vram:
+            self.image_cond_model.to(self.device)
+        cond = self.image_cond_model(image)
+        if self.low_vram:
+            self.image_cond_model.cpu()
+        if not include_neg_cond:
+            return {'cond': cond}
+        neg_cond = torch.zeros_like(cond)
+        return {
+            'cond': cond,
+            'neg_cond': neg_cond,
+        }
+    
+    def encode_shape_slat(
+        self,
+        mesh: trimesh.Trimesh,
+        resolution: int = 1024,
+    ) -> SparseTensor:
+        """
+        Encode the meshes to structured latent.
+
+        Args:
+            mesh (trimesh.Trimesh): The mesh to encode.
+            resolution (int): The resolution of mesh
+        
+        Returns:
+            SparseTensor: The encoded structured latent.
+        """
+        vertices = torch.from_numpy(mesh.vertices).float()
+        faces = torch.from_numpy(mesh.faces).long()
+        
+        voxel_indices, dual_vertices, intersected = o_voxel.convert.mesh_to_flexible_dual_grid(
+            vertices.cpu(), faces.cpu(),
+            grid_size=resolution,
+            aabb=[[-0.5,-0.5,-0.5],[0.5,0.5,0.5]],
+            face_weight=1.0,
+            boundary_weight=0.2,
+            regularization_weight=1e-2,
+            timing=True,
+        )
+            
+        vertices = SparseTensor(
+            feats=dual_vertices * resolution - voxel_indices,
+            coords=torch.cat([torch.zeros_like(voxel_indices[:, 0:1]), voxel_indices], dim=-1)
+        ).to(self.device)
+        intersected = vertices.replace(intersected).to(self.device)
+            
+        if self.low_vram:
+            self.models['shape_slat_encoder'].to(self.device)
+        shape_slat = self.models['shape_slat_encoder'](vertices, intersected)
+        if self.low_vram:
+            self.models['shape_slat_encoder'].cpu()
+        return shape_slat
+
+    def sample_tex_slat(
+        self,
+        cond: dict,
+        flow_model,
+        shape_slat: SparseTensor,
+        sampler_params: dict = {},
+    ) -> SparseTensor:
+        """
+        Sample structured latent with the given conditioning.
+        
+        Args:
+            cond (dict): The conditioning information.
+            shape_slat (SparseTensor): The structured latent for shape
+            sampler_params (dict): Additional parameters for the sampler.
+        """
+        # Sample structured latent
+        std = torch.tensor(self.shape_slat_normalization['std'])[None].to(shape_slat.device)
+        mean = torch.tensor(self.shape_slat_normalization['mean'])[None].to(shape_slat.device)
+        shape_slat = (shape_slat - mean) / std
+
+        in_channels = flow_model.in_channels if isinstance(flow_model, nn.Module) else flow_model[0].in_channels
+        noise = shape_slat.replace(feats=torch.randn(shape_slat.coords.shape[0], in_channels - shape_slat.feats.shape[1]).to(self.device))
+        sampler_params = {**self.tex_slat_sampler_params, **sampler_params}
+        if self.low_vram:
+            flow_model.to(self.device)
+        slat = self.tex_slat_sampler.sample(
+            flow_model,
+            noise,
+            concat_cond=shape_slat,
+            **cond,
+            **sampler_params,
+            verbose=True,
+            tqdm_desc="Sampling texture SLat",
+        ).samples
+        if self.low_vram:
+            flow_model.cpu()
+
+        std = torch.tensor(self.tex_slat_normalization['std'])[None].to(slat.device)
+        mean = torch.tensor(self.tex_slat_normalization['mean'])[None].to(slat.device)
+        slat = slat * std + mean
+        
+        return slat
+
+    def decode_tex_slat(
+        self,
+        slat: SparseTensor,
+    ) -> SparseTensor:
+        """
+        Decode the structured latent.
+
+        Args:
+            slat (SparseTensor): The structured latent.
+
+        Returns:
+            SparseTensor: The decoded texture voxels
+        """
+        if self.low_vram:
+            self.models['tex_slat_decoder'].to(self.device)
+        ret = self.models['tex_slat_decoder'](slat) * 0.5 + 0.5
+        if self.low_vram:
+            self.models['tex_slat_decoder'].cpu()
+        return ret
+    
+    def postprocess_mesh(
+        self,
+        mesh: trimesh.Trimesh,
+        pbr_voxel: SparseTensor,
+        resolution: int = 1024,
+        texture_size: int = 1024,
+    ) -> trimesh.Trimesh:
+        vertices = mesh.vertices
+        faces = mesh.faces
+        normals = mesh.vertex_normals
+        vertices_torch = torch.from_numpy(vertices).float().cuda()
+        faces_torch = torch.from_numpy(faces).int().cuda()
+        if hasattr(mesh, 'visual') and hasattr(mesh.visual, 'uv') and mesh.visual.uv is not None:
+            uvs = mesh.visual.uv.copy()
+            uvs[:, 1] = 1 - uvs[:, 1]
+            uvs_torch = torch.from_numpy(uvs).float().cuda()
+        else:
+            _cumesh = cumesh.CuMesh()
+            _cumesh.init(vertices_torch, faces_torch)
+            vertices_torch, faces_torch, uvs_torch, vmap = _cumesh.uv_unwrap(return_vmaps=True)
+            vertices_torch = vertices_torch.cuda()
+            faces_torch = faces_torch.cuda()
+            uvs_torch = uvs_torch.cuda()
+            vertices = vertices_torch.cpu().numpy()
+            faces = faces_torch.cpu().numpy()
+            uvs = uvs_torch.cpu().numpy()
+            normals = normals[vmap.cpu().numpy()]
+                
+        # rasterize
+        ctx = dr.RasterizeCudaContext()
+        uvs_torch = torch.cat([uvs_torch * 2 - 1, torch.zeros_like(uvs_torch[:, :1]), torch.ones_like(uvs_torch[:, :1])], dim=-1).unsqueeze(0)
+        rast, _ = dr.rasterize(
+            ctx, uvs_torch, faces_torch,
+            resolution=[texture_size, texture_size],
+        )
+        mask = rast[0, ..., 3] > 0
+        pos = dr.interpolate(vertices_torch.unsqueeze(0), rast, faces_torch)[0][0]
+        
+        attrs = torch.zeros(texture_size, texture_size, pbr_voxel.shape[1], device=self.device)
+        attrs[mask] = flex_gemm.ops.grid_sample.grid_sample_3d(
+            pbr_voxel.feats,
+            pbr_voxel.coords,
+            shape=torch.Size([*pbr_voxel.shape, *pbr_voxel.spatial_shape]),
+            grid=((pos[mask] + 0.5) * resolution).reshape(1, -1, 3),
+            mode='trilinear',
+        )
+        
+        # construct mesh
+        mask = mask.cpu().numpy()
+        base_color = np.clip(attrs[..., self.pbr_attr_layout['base_color']].cpu().numpy() * 255, 0, 255).astype(np.uint8)
+        metallic = np.clip(attrs[..., self.pbr_attr_layout['metallic']].cpu().numpy() * 255, 0, 255).astype(np.uint8)
+        roughness = np.clip(attrs[..., self.pbr_attr_layout['roughness']].cpu().numpy() * 255, 0, 255).astype(np.uint8)
+        alpha = np.clip(attrs[..., self.pbr_attr_layout['alpha']].cpu().numpy() * 255, 0, 255).astype(np.uint8)
+        
+        # extend
+        mask = (~mask).astype(np.uint8)
+        base_color = cv2.inpaint(base_color, mask, 3, cv2.INPAINT_TELEA)
+        metallic = cv2.inpaint(metallic, mask, 1, cv2.INPAINT_TELEA)[..., None]
+        roughness = cv2.inpaint(roughness, mask, 1, cv2.INPAINT_TELEA)[..., None]
+        alpha = cv2.inpaint(alpha, mask, 1, cv2.INPAINT_TELEA)[..., None]
+        
+        material = trimesh.visual.material.PBRMaterial(
+            baseColorTexture=Image.fromarray(np.concatenate([base_color, alpha], axis=-1)),
+            baseColorFactor=np.array([255, 255, 255, 255], dtype=np.uint8),
+            metallicRoughnessTexture=Image.fromarray(np.concatenate([np.zeros_like(metallic), roughness, metallic], axis=-1)),
+            metallicFactor=1.0,
+            roughnessFactor=1.0,
+            alphaMode='OPAQUE',
+            doubleSided=True,
+        )
+
+        # Swap Y and Z axes, invert Y (common conversion for GLB compatibility)
+        vertices[:, 1], vertices[:, 2] = vertices[:, 2], -vertices[:, 1]
+        normals[:, 1], normals[:, 2] = normals[:, 2], -normals[:, 1]
+        uvs[:, 1] = 1 - uvs[:, 1] # Flip UV V-coordinate
+        
+        textured_mesh = trimesh.Trimesh(
+            vertices=vertices,
+            faces=faces,
+            vertex_normals=normals,
+            process=False,
+            visual=trimesh.visual.TextureVisuals(uv=uvs, material=material)
+        )
+        
+        return textured_mesh
+        
+    
+    @torch.no_grad()
+    def run(
+        self,
+        mesh: trimesh.Trimesh,
+        image: Image.Image,
+        seed: int = 42,
+        tex_slat_sampler_params: dict = {},
+        preprocess_image: bool = True,
+        resolution: int = 1024,
+        texture_size: int = 2048,
+    ) -> trimesh.Trimesh:
+        """
+        Run the pipeline.
+
+        Args:
+            mesh (trimesh.Trimesh): The mesh to texture.
+            image (Image.Image): The image prompt.
+            seed (int): The random seed.
+            tex_slat_sampler_params (dict): Additional parameters for the texture latent sampler.
+            preprocess_image (bool): Whether to preprocess the image.
+        """
+        if preprocess_image:
+            image = self.preprocess_image(image)
+        mesh = self.preprocess_mesh(mesh)
+        torch.manual_seed(seed)
+        cond = self.get_cond([image], 512) if resolution == 512 else self.get_cond([image], 1024)
+        shape_slat = self.encode_shape_slat(mesh, resolution)
+        tex_model = self.models['tex_slat_flow_model_512'] if resolution == 512 else self.models['tex_slat_flow_model_1024']
+        tex_slat = self.sample_tex_slat(
+            cond, tex_model,
+            shape_slat, tex_slat_sampler_params
+        )
+        pbr_voxel = self.decode_tex_slat(tex_slat)
+        out_mesh = self.postprocess_mesh(mesh, pbr_voxel, resolution, texture_size)
+        return out_mesh
diff --git a/pixal3d/renderers/__init__.py b/pixal3d/renderers/__init__.py
new file mode 100644
index 0000000..de3203d
--- /dev/null
+++ b/pixal3d/renderers/__init__.py
@@ -0,0 +1,33 @@
+import importlib
+
+__attributes = {
+    'MeshRenderer': 'mesh_renderer',
+    'VoxelRenderer': 'voxel_renderer',
+    'PbrMeshRenderer': 'pbr_mesh_renderer',
+    'EnvMap': 'pbr_mesh_renderer',
+}
+
+__submodules = []
+
+__all__ = list(__attributes.keys()) + __submodules
+
+def __getattr__(name):
+    if name not in globals():
+        if name in __attributes:
+            module_name = __attributes[name]
+            module = importlib.import_module(f".{module_name}", __name__)
+            globals()[name] = getattr(module, name)
+        elif name in __submodules:
+            module = importlib.import_module(f".{name}", __name__)
+            globals()[name] = module
+        else:
+            raise AttributeError(f"module {__name__} has no attribute {name}")
+    return globals()[name]
+
+
+# For Pylance
+if __name__ == '__main__':
+    from .mesh_renderer import MeshRenderer
+    from .voxel_renderer import VoxelRenderer
+    from .pbr_mesh_renderer import PbrMeshRenderer, EnvMap
+    
\ No newline at end of file
diff --git a/pixal3d/renderers/mesh_renderer.py b/pixal3d/renderers/mesh_renderer.py
new file mode 100644
index 0000000..e20efc5
--- /dev/null
+++ b/pixal3d/renderers/mesh_renderer.py
@@ -0,0 +1,414 @@
+from typing import *
+import torch
+from easydict import EasyDict as edict
+from ..representations.mesh import Mesh, MeshWithVoxel, MeshWithPbrMaterial, TextureFilterMode, AlphaMode, TextureWrapMode
+import torch.nn.functional as F
+
+
+def intrinsics_to_projection(
+        intrinsics: torch.Tensor,
+        near: float,
+        far: float,
+    ) -> torch.Tensor:
+    """
+    OpenCV intrinsics to OpenGL perspective matrix
+
+    Args:
+        intrinsics (torch.Tensor): [3, 3] OpenCV intrinsics matrix
+        near (float): near plane to clip
+        far (float): far plane to clip
+    Returns:
+        (torch.Tensor): [4, 4] OpenGL perspective matrix
+    """
+    fx, fy = intrinsics[0, 0], intrinsics[1, 1]
+    cx, cy = intrinsics[0, 2], intrinsics[1, 2]
+    ret = torch.zeros((4, 4), dtype=intrinsics.dtype, device=intrinsics.device)
+    ret[0, 0] = 2 * fx
+    ret[1, 1] = 2 * fy
+    ret[0, 2] = 2 * cx - 1
+    ret[1, 2] = - 2 * cy + 1
+    ret[2, 2] = (far + near) / (far - near)
+    ret[2, 3] = 2 * near * far / (near - far)
+    ret[3, 2] = 1.
+    return ret
+    
+
+class MeshRenderer:
+    """
+    Renderer for the Mesh representation.
+
+    Args:
+        rendering_options (dict): Rendering options.
+        """
+    def __init__(self, rendering_options={}, device='cuda'):
+        if 'dr' not in globals():
+            import nvdiffrast.torch as dr
+        
+        self.rendering_options = edict({
+            "resolution": None,
+            "near": None,
+            "far": None,
+            "ssaa": 1,
+            "chunk_size": None,
+            "antialias": True,
+            "clamp_barycentric_coords": False,
+        })
+        self.rendering_options.update(rendering_options)
+        self.glctx = dr.RasterizeCudaContext(device=device)
+        self.device=device
+        
+    def render(
+            self,
+            mesh : Mesh,
+            extrinsics: torch.Tensor,
+            intrinsics: torch.Tensor,
+            return_types = ["mask", "normal", "depth"],
+            transformation : Optional[torch.Tensor] = None
+        ) -> edict:
+        """
+        Render the mesh.
+
+        Args:
+            mesh : meshmodel
+            extrinsics (torch.Tensor): (4, 4) camera extrinsics
+            intrinsics (torch.Tensor): (3, 3) camera intrinsics
+            return_types (list): list of return types, can be "attr", "mask", "depth", "coord", "normal"
+
+        Returns:
+            edict based on return_types containing:
+                attr (torch.Tensor): [C, H, W] rendered attr image
+                depth (torch.Tensor): [H, W] rendered depth image
+                normal (torch.Tensor): [3, H, W] rendered normal image
+                mask (torch.Tensor): [H, W] rendered mask image
+        """
+        if 'dr' not in globals():
+            import nvdiffrast.torch as dr
+            
+        resolution = self.rendering_options["resolution"]
+        near = self.rendering_options["near"]
+        far = self.rendering_options["far"]
+        ssaa = self.rendering_options["ssaa"]
+        chunk_size = self.rendering_options["chunk_size"]
+        antialias = self.rendering_options["antialias"]
+        clamp_barycentric_coords = self.rendering_options["clamp_barycentric_coords"]
+        
+        if mesh.vertices.shape[0] == 0 or mesh.faces.shape[0] == 0:
+            ret_dict = edict()
+            for type in return_types:
+                if type == "mask" :
+                    ret_dict[type] = torch.zeros((resolution, resolution), dtype=torch.float32, device=self.device)
+                elif type == "depth":
+                    ret_dict[type] = torch.zeros((resolution, resolution), dtype=torch.float32, device=self.device)
+                elif type == "normal":
+                    ret_dict[type] = torch.full((3, resolution, resolution), 0.5, dtype=torch.float32, device=self.device)
+                elif type == "coord":
+                    ret_dict[type] = torch.zeros((3, resolution, resolution), dtype=torch.float32, device=self.device)
+                elif type == "attr":
+                    if isinstance(mesh, MeshWithVoxel):
+                        ret_dict[type] = torch.zeros((mesh.attrs.shape[-1], resolution, resolution), dtype=torch.float32, device=self.device)
+                    else:
+                        ret_dict[type] = torch.zeros((mesh.vertex_attrs.shape[-1], resolution, resolution), dtype=torch.float32, device=self.device)
+            return ret_dict
+        
+        perspective = intrinsics_to_projection(intrinsics, near, far)
+        
+        full_proj = (perspective @ extrinsics).unsqueeze(0)
+        extrinsics = extrinsics.unsqueeze(0)
+        
+        vertices = mesh.vertices.unsqueeze(0)
+        vertices_homo = torch.cat([vertices, torch.ones_like(vertices[..., :1])], dim=-1)
+        if transformation is not None:
+            vertices_homo = torch.bmm(vertices_homo, transformation.unsqueeze(0).transpose(-1, -2))
+            vertices = vertices_homo[..., :3].contiguous()
+        vertices_camera = torch.bmm(vertices_homo, extrinsics.transpose(-1, -2))
+        vertices_clip = torch.bmm(vertices_homo, full_proj.transpose(-1, -2))
+        faces = mesh.faces
+        
+        if 'normal' in return_types:
+            v0 = vertices_camera[0, mesh.faces[:, 0], :3]
+            v1 = vertices_camera[0, mesh.faces[:, 1], :3]
+            v2 = vertices_camera[0, mesh.faces[:, 2], :3]
+            e0 = v1 - v0
+            e1 = v2 - v0
+            face_normal = torch.cross(e0, e1, dim=1)
+            face_normal = F.normalize(face_normal, dim=1)
+            face_normal = torch.where(torch.sum(face_normal * v0, dim=1, keepdim=True) > 0, face_normal, -face_normal)
+        
+        out_dict = edict()
+        if chunk_size is None:
+            rast, rast_db = dr.rasterize(
+                self.glctx, vertices_clip, faces, (resolution * ssaa, resolution * ssaa)
+            )
+            if clamp_barycentric_coords:
+                rast[..., :2] = torch.clamp(rast[..., :2], 0, 1)
+                rast[..., :2] /= torch.where(rast[..., :2].sum(dim=-1, keepdim=True) > 1, rast[..., :2].sum(dim=-1, keepdim=True), torch.ones_like(rast[..., :2]))
+            for type in return_types:
+                img = None
+                if type == "mask" :
+                    img = (rast[..., -1:] > 0).float()
+                    if antialias: img = dr.antialias(img, rast, vertices_clip, faces)
+                elif type == "depth":
+                    img = dr.interpolate(vertices_camera[..., 2:3].contiguous(), rast, faces)[0]
+                    if antialias: img = dr.antialias(img, rast, vertices_clip, faces)
+                elif type == "normal" :
+                    img = dr.interpolate(face_normal.unsqueeze(0), rast, torch.arange(face_normal.shape[0], dtype=torch.int, device=self.device).unsqueeze(1).repeat(1, 3).contiguous())[0]
+                    if antialias: img = dr.antialias(img, rast, vertices_clip, faces)
+                    img = (img + 1) / 2
+                elif type == "coord":
+                    img = dr.interpolate(vertices, rast, faces)[0]
+                    if antialias: img = dr.antialias(img, rast, vertices_clip, faces)
+                elif type == "attr":
+                    if isinstance(mesh, MeshWithVoxel):
+                        if 'grid_sample_3d' not in globals():
+                            from flex_gemm.ops.grid_sample import grid_sample_3d
+                        mask = rast[..., -1:] > 0
+                        xyz = dr.interpolate(vertices, rast, faces)[0]
+                        xyz = ((xyz - mesh.origin) / mesh.voxel_size).reshape(1, -1, 3)
+                        img = grid_sample_3d(
+                            mesh.attrs,
+                            torch.cat([torch.zeros_like(mesh.coords[..., :1]), mesh.coords], dim=-1),
+                            mesh.voxel_shape,
+                            xyz,
+                            mode='trilinear'
+                        )
+                        img = img.reshape(1, resolution * ssaa, resolution * ssaa, mesh.attrs.shape[-1]) * mask
+                    elif isinstance(mesh, MeshWithPbrMaterial):
+                        tri_id = rast[0, :, :, -1:]
+                        mask = tri_id > 0
+                        uv_coords = mesh.uv_coords.reshape(1, -1, 2)
+                        texc, texd = dr.interpolate(
+                            uv_coords,
+                            rast,
+                            torch.arange(mesh.uv_coords.shape[0] * 3, dtype=torch.int, device=self.device).reshape(-1, 3),
+                            rast_db=rast_db,
+                            diff_attrs='all'
+                        )
+                        # Fix problematic texture coordinates
+                        texc = torch.nan_to_num(texc, nan=0.0, posinf=1e3, neginf=-1e3)
+                        texc = torch.clamp(texc, min=-1e3, max=1e3)
+                        texd = torch.nan_to_num(texd, nan=0.0, posinf=1e3, neginf=-1e3)
+                        texd = torch.clamp(texd, min=-1e3, max=1e3)
+                        mid = mesh.material_ids[(tri_id - 1).long()]
+                        imgs = {
+                            'base_color': torch.zeros((resolution * ssaa, resolution * ssaa, 3), dtype=torch.float32, device=self.device),
+                            'metallic': torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device),
+                            'roughness': torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device),
+                            'alpha': torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device)
+                        }
+                        for id, mat in enumerate(mesh.materials):
+                            mat_mask = (mid == id).float() * mask.float()
+                            mat_texc = texc * mat_mask
+                            mat_texd = texd * mat_mask
+
+                            if mat.base_color_texture is not None:
+                                base_color = dr.texture(
+                                    mat.base_color_texture.image.unsqueeze(0),
+                                    mat_texc,
+                                    mat_texd,
+                                    filter_mode='linear-mipmap-linear' if mat.base_color_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest',
+                                    boundary_mode='clamp' if mat.base_color_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap'
+                                )[0]
+                                imgs['base_color'] += base_color * mat.base_color_factor * mat_mask
+                            else:
+                                imgs['base_color'] += mat.base_color_factor * mat_mask
+                                
+                            if mat.metallic_texture is not None:
+                                metallic = dr.texture(
+                                    mat.metallic_texture.image.unsqueeze(0),
+                                    mat_texc,
+                                    mat_texd,
+                                    filter_mode='linear-mipmap-linear' if mat.metallic_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest',
+                                    boundary_mode='clamp' if mat.metallic_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap'
+                                )[0]
+                                imgs['metallic'] += metallic * mat.metallic_factor * mat_mask
+                            else:
+                                imgs['metallic'] += mat.metallic_factor * mat_mask
+
+                            if mat.roughness_texture is not None:
+                                roughness = dr.texture(
+                                    mat.roughness_texture.image.unsqueeze(0),
+                                    mat_texc,
+                                    mat_texd,
+                                    filter_mode='linear-mipmap-linear' if mat.roughness_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest',
+                                    boundary_mode='clamp' if mat.roughness_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap'
+                                )[0]
+                                imgs['roughness'] += roughness * mat.roughness_factor * mat_mask
+                            else:
+                                imgs['roughness'] += mat.roughness_factor * mat_mask
+
+                            if mat.alpha_mode == AlphaMode.OPAQUE:
+                                imgs['alpha'] += 1.0 * mat_mask
+                            else:
+                                if mat.alpha_texture is not None:
+                                    alpha = dr.texture(
+                                        mat.alpha_texture.image.unsqueeze(0),
+                                        mat_texc,
+                                        mat_texd,
+                                        filter_mode='linear-mipmap-linear' if mat.alpha_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest',
+                                        boundary_mode='clamp' if mat.alpha_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap'
+                                    )[0]
+                                    if mat.alpha_mode == AlphaMode.MASK:
+                                        imgs['alpha'] += (alpha * mat.alpha_factor > mat.alpha_cutoff).float() * mat_mask
+                                    elif mat.alpha_mode == AlphaMode.BLEND:
+                                        imgs['alpha'] += alpha * mat.alpha_factor * mat_mask
+                                else:
+                                    if mat.alpha_mode == AlphaMode.MASK:
+                                        imgs['alpha'] += (mat.alpha_factor > mat.alpha_cutoff).float() * mat_mask
+                                    elif mat.alpha_mode == AlphaMode.BLEND:
+                                        imgs['alpha'] += mat.alpha_factor * mat_mask
+                    
+                        img = torch.cat([imgs[name] for name in imgs.keys()], dim=-1).unsqueeze(0)
+                    else:
+                        img = dr.interpolate(mesh.vertex_attrs.unsqueeze(0), rast, faces)[0]
+                        if antialias: img = dr.antialias(img, rast, vertices_clip, faces)
+                        
+                out_dict[type] = img
+        else:
+            z_buffer = torch.full((1, resolution * ssaa, resolution * ssaa), torch.inf, device=self.device, dtype=torch.float32)
+            for i in range(0, faces.shape[0], chunk_size):
+                faces_chunk = faces[i:i+chunk_size]
+                rast, rast_db = dr.rasterize(
+                    self.glctx, vertices_clip, faces_chunk, (resolution * ssaa, resolution * ssaa)
+                )
+                z_filter = torch.logical_and(
+                    rast[..., 3] != 0,
+                    rast[..., 2] < z_buffer
+                )
+                z_buffer[z_filter] = rast[z_filter][..., 2]
+            
+                for type in return_types:
+                    img = None
+                    if type == "mask" :
+                        img = (rast[..., -1:] > 0).float()
+                    elif type == "depth":
+                        img = dr.interpolate(vertices_camera[..., 2:3].contiguous(), rast, faces_chunk)[0]
+                    elif type == "normal" :
+                        face_normal_chunk = face_normal[i:i+chunk_size]
+                        img = dr.interpolate(face_normal_chunk.unsqueeze(0), rast, torch.arange(face_normal_chunk.shape[0], dtype=torch.int, device=self.device).unsqueeze(1).repeat(1, 3).contiguous())[0]
+                        img = (img + 1) / 2
+                    elif type == "coord":
+                        img = dr.interpolate(vertices, rast, faces_chunk)[0]
+                    elif type == "attr":
+                        if isinstance(mesh, MeshWithVoxel):
+                            if 'grid_sample_3d' not in globals():
+                                from flex_gemm.ops.grid_sample import grid_sample_3d
+                            mask = rast[..., -1:] > 0
+                            xyz = dr.interpolate(vertices, rast, faces_chunk)[0]
+                            xyz = ((xyz - mesh.origin) / mesh.voxel_size).reshape(1, -1, 3)
+                            img = grid_sample_3d(
+                                mesh.attrs,
+                                torch.cat([torch.zeros_like(mesh.coords[..., :1]), mesh.coords], dim=-1),
+                                mesh.voxel_shape,
+                                xyz,
+                                mode='trilinear'
+                            )
+                            img = img.reshape(1, resolution * ssaa, resolution * ssaa, mesh.attrs.shape[-1]) * mask
+                        elif isinstance(mesh, MeshWithPbrMaterial):
+                            tri_id = rast[0, :, :, -1:]
+                            mask = tri_id > 0
+                            uv_coords = mesh.uv_coords.reshape(1, -1, 2)
+                            texc, texd = dr.interpolate(
+                                uv_coords,
+                                rast,
+                                torch.arange(mesh.uv_coords.shape[0] * 3, dtype=torch.int, device=self.device).reshape(-1, 3),
+                                rast_db=rast_db,
+                                diff_attrs='all'
+                            )
+                            # Fix problematic texture coordinates
+                            texc = torch.nan_to_num(texc, nan=0.0, posinf=1e3, neginf=-1e3)
+                            texc = torch.clamp(texc, min=-1e3, max=1e3)
+                            texd = torch.nan_to_num(texd, nan=0.0, posinf=1e3, neginf=-1e3)
+                            texd = torch.clamp(texd, min=-1e3, max=1e3)
+                            mid = mesh.material_ids[(tri_id - 1).long()]
+                            imgs = {
+                                'base_color': torch.zeros((resolution * ssaa, resolution * ssaa, 3), dtype=torch.float32, device=self.device),
+                                'metallic': torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device),
+                                'roughness': torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device),
+                                'alpha': torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device)
+                            }
+                            for id, mat in enumerate(mesh.materials):
+                                mat_mask = (mid == id).float() * mask.float()
+                                mat_texc = texc * mat_mask
+                                mat_texd = texd * mat_mask
+
+                                if mat.base_color_texture is not None:
+                                    base_color = dr.texture(
+                                        mat.base_color_texture.image.unsqueeze(0),
+                                        mat_texc,
+                                        mat_texd,
+                                        filter_mode='linear-mipmap-linear' if mat.base_color_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest',
+                                        boundary_mode='clamp' if mat.base_color_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap'
+                                    )[0]
+                                    imgs['base_color'] += base_color * mat.base_color_factor * mat_mask
+                                else:
+                                    imgs['base_color'] += mat.base_color_factor * mat_mask
+                                    
+                                if mat.metallic_texture is not None:
+                                    metallic = dr.texture(
+                                        mat.metallic_texture.image.unsqueeze(0),
+                                        mat_texc,
+                                        mat_texd,
+                                        filter_mode='linear-mipmap-linear' if mat.metallic_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest',
+                                        boundary_mode='clamp' if mat.metallic_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap'
+                                    )[0]
+                                    imgs['metallic'] += metallic * mat.metallic_factor * mat_mask
+                                else:
+                                    imgs['metallic'] += mat.metallic_factor * mat_mask
+
+                                if mat.roughness_texture is not None:
+                                    roughness = dr.texture(
+                                        mat.roughness_texture.image.unsqueeze(0),
+                                        mat_texc,
+                                        mat_texd,
+                                        filter_mode='linear-mipmap-linear' if mat.roughness_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest',
+                                        boundary_mode='clamp' if mat.roughness_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap'
+                                    )[0]
+                                    imgs['roughness'] += roughness * mat.roughness_factor * mat_mask
+                                else:
+                                    imgs['roughness'] += mat.roughness_factor * mat_mask
+
+                                if mat.alpha_mode == AlphaMode.OPAQUE:
+                                    imgs['alpha'] += 1.0 * mat_mask
+                                else:
+                                    if mat.alpha_texture is not None:
+                                        alpha = dr.texture(
+                                            mat.alpha_texture.image.unsqueeze(0),
+                                            mat_texc,
+                                            mat_texd,
+                                            filter_mode='linear-mipmap-linear' if mat.alpha_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest',
+                                            boundary_mode='clamp' if mat.alpha_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap'
+                                        )[0]
+                                        if mat.alpha_mode == AlphaMode.MASK:
+                                            imgs['alpha'] += (alpha * mat.alpha_factor > mat.alpha_cutoff).float() * mat_mask
+                                        elif mat.alpha_mode == AlphaMode.BLEND:
+                                            imgs['alpha'] += alpha * mat.alpha_factor * mat_mask
+                                    else:
+                                        if mat.alpha_mode == AlphaMode.MASK:
+                                            imgs['alpha'] += (mat.alpha_factor > mat.alpha_cutoff).float() * mat_mask
+                                        elif mat.alpha_mode == AlphaMode.BLEND:
+                                            imgs['alpha'] += mat.alpha_factor * mat_mask
+                        
+                            img = torch.cat([imgs[name] for name in imgs.keys()], dim=-1).unsqueeze(0)
+                        else:
+                            img = dr.interpolate(mesh.vertex_attrs.unsqueeze(0), rast, faces_chunk)[0]
+                            
+                    if type not in out_dict:
+                        out_dict[type] = img
+                    else:
+                        out_dict[type][z_filter] = img[z_filter]
+
+        for type in return_types:
+            img = out_dict[type]
+            if ssaa > 1:
+                img = F.interpolate(img.permute(0, 3, 1, 2), (resolution, resolution), mode='bilinear', align_corners=False, antialias=True)
+                img = img.squeeze()
+            else:
+                img = img.permute(0, 3, 1, 2).squeeze()
+            out_dict[type] = img
+
+        if isinstance(mesh, (MeshWithVoxel, MeshWithPbrMaterial)) and 'attr' in return_types:
+            for k, s in mesh.layout.items():
+                out_dict[k] = out_dict['attr'][s]
+            del out_dict['attr']
+        
+        return out_dict
diff --git a/pixal3d/renderers/pbr_mesh_renderer.py b/pixal3d/renderers/pbr_mesh_renderer.py
new file mode 100644
index 0000000..14ded30
--- /dev/null
+++ b/pixal3d/renderers/pbr_mesh_renderer.py
@@ -0,0 +1,521 @@
+from typing import *
+import torch
+from easydict import EasyDict as edict
+import numpy as np
+import utils3d
+from ..representations.mesh import Mesh, MeshWithVoxel, MeshWithPbrMaterial, TextureFilterMode, AlphaMode, TextureWrapMode
+import torch.nn.functional as F
+
+
+def cube_to_dir(s, x, y):
+    if s == 0:   rx, ry, rz = torch.ones_like(x), -x, -y
+    elif s == 1: rx, ry, rz = -torch.ones_like(x), x, -y
+    elif s == 2: rx, ry, rz = x, y, torch.ones_like(x)
+    elif s == 3: rx, ry, rz = x, -y, -torch.ones_like(x)
+    elif s == 4: rx, ry, rz = x, torch.ones_like(x), -y
+    elif s == 5: rx, ry, rz = -x, -torch.ones_like(x), -y
+    return torch.stack((rx, ry, rz), dim=-1)
+
+
+def latlong_to_cubemap(latlong_map, res):
+    if 'dr' not in globals():
+        import nvdiffrast.torch as dr
+    cubemap = torch.zeros(6, res[0], res[1], latlong_map.shape[-1], dtype=torch.float32, device='cuda')
+    for s in range(6):
+        gy, gx = torch.meshgrid(torch.linspace(-1.0 + 1.0 / res[0], 1.0 - 1.0 / res[0], res[0], device='cuda'), 
+                                torch.linspace(-1.0 + 1.0 / res[1], 1.0 - 1.0 / res[1], res[1], device='cuda'),
+                                indexing='ij')
+        v = F.normalize(cube_to_dir(s, gx, gy), dim=-1)
+
+        tu = torch.atan2(v[..., 0:1], -v[..., 2:3]) / (2 * np.pi) + 0.5
+        tv = torch.acos(torch.clamp(v[..., 1:2], min=-1, max=1)) / np.pi
+        texcoord = torch.cat((tu, tv), dim=-1)
+
+        cubemap[s, ...] = dr.texture(latlong_map[None, ...], texcoord[None, ...], filter_mode='linear')[0]
+    return cubemap
+
+
+class EnvMap:
+    def __init__(self, image: torch.Tensor):
+        self.image = image
+        
+    @property
+    def _backend(self):
+        if not hasattr(self, '_nvdiffrec_envlight'):
+            if 'EnvironmentLight' not in globals():
+                from nvdiffrec_render.light import EnvironmentLight
+            cubemap = latlong_to_cubemap(self.image, [512, 512])
+            self._nvdiffrec_envlight = EnvironmentLight(cubemap)
+            self._nvdiffrec_envlight.build_mips()
+        return self._nvdiffrec_envlight
+
+    def shade(self, gb_pos, gb_normal, kd, ks, view_pos, specular=True):
+        return self._backend.shade(gb_pos, gb_normal, kd, ks, view_pos, specular)
+    
+    def sample(self, directions: torch.Tensor):
+        if 'dr' not in globals():
+            import nvdiffrast.torch as dr
+        return dr.texture(
+            self._backend.base.unsqueeze(0),
+            directions.unsqueeze(0),
+            boundary_mode='cube',
+        )[0]
+            
+
+def intrinsics_to_projection(
+        intrinsics: torch.Tensor,
+        near: float,
+        far: float,
+    ) -> torch.Tensor:
+    """
+    OpenCV intrinsics to OpenGL perspective matrix
+
+    Args:
+        intrinsics (torch.Tensor): [3, 3] OpenCV intrinsics matrix
+        near (float): near plane to clip
+        far (float): far plane to clip
+    Returns:
+        (torch.Tensor): [4, 4] OpenGL perspective matrix
+    """
+    fx, fy = intrinsics[0, 0], intrinsics[1, 1]
+    cx, cy = intrinsics[0, 2], intrinsics[1, 2]
+    ret = torch.zeros((4, 4), dtype=intrinsics.dtype, device=intrinsics.device)
+    ret[0, 0] = 2 * fx
+    ret[1, 1] = 2 * fy
+    ret[0, 2] = 2 * cx - 1
+    ret[1, 2] = - 2 * cy + 1
+    ret[2, 2] = (far + near) / (far - near)
+    ret[2, 3] = 2 * near * far / (near - far)
+    ret[3, 2] = 1.
+    return ret
+
+
+def screen_space_ambient_occlusion(
+    depth: torch.Tensor,
+    normal: torch.Tensor,
+    perspective: torch.Tensor,
+    radius: float = 0.1,
+    bias: float = 1e-6,
+    samples: int = 64,
+    intensity: float = 1.0,
+) -> torch.Tensor:
+    """
+    Screen space ambient occlusion (SSAO)
+
+    Args:
+        depth (torch.Tensor): [H, W, 1] depth image
+        normal (torch.Tensor): [H, W, 3] normal image
+        perspective (torch.Tensor): [4, 4] camera projection matrix
+        radius (float): radius of the SSAO kernel
+        bias (float): bias to avoid self-occlusion
+        samples (int): number of samples to use for the SSAO kernel
+        intensity (float): intensity of the SSAO effect
+    Returns:
+        (torch.Tensor): [H, W, 1] SSAO image
+    """
+    device = depth.device
+    H, W, _ = depth.shape
+    
+    fx = perspective[0, 0]
+    fy = perspective[1, 1]
+    cx = perspective[0, 2]
+    cy = perspective[1, 2]
+    
+    y_grid, x_grid = torch.meshgrid(
+        (torch.arange(H, device=device) + 0.5) / H * 2 - 1,
+        (torch.arange(W, device=device) + 0.5) / W * 2 - 1,
+        indexing='ij'
+    )
+    x_view = (x_grid.float() - cx) * depth[..., 0] / fx
+    y_view = (y_grid.float() - cy) * depth[..., 0] / fy
+    view_pos = torch.stack([x_view, y_view, depth[..., 0]], dim=-1) # [H, W, 3]
+    
+    depth_feat = depth.permute(2, 0, 1).unsqueeze(0)
+    occlusion = torch.zeros((H, W), device=device)
+    
+    # start sampling
+    for _ in range(samples):
+        # sample normal distribution, if inside, flip the sign
+        rnd_vec = torch.randn(H, W, 3, device=device)
+        rnd_vec = F.normalize(rnd_vec, p=2, dim=-1)
+        dot_val = torch.sum(rnd_vec * normal, dim=-1, keepdim=True)
+        sample_dir = torch.sign(dot_val) * rnd_vec
+        scale = torch.rand(H, W, 1, device=device)
+        scale = scale * scale
+        sample_pos = view_pos + sample_dir * radius * scale
+        sample_z = sample_pos[..., 2]
+        
+        # project to screen space
+        z_safe = torch.clamp(sample_pos[..., 2], min=1e-5)
+        proj_u = (sample_pos[..., 0] * fx / z_safe) + cx
+        proj_v = (sample_pos[..., 1] * fy / z_safe) + cy
+        grid = torch.stack([proj_u, proj_v], dim=-1).unsqueeze(0)
+        geo_z = F.grid_sample(depth_feat, grid, mode='nearest', padding_mode='border').squeeze()
+        range_check = torch.abs(geo_z - sample_z) < radius
+        is_occluded = (geo_z <= sample_z - bias) & range_check
+        occlusion += is_occluded.float()
+        
+    f_occ = occlusion / samples * intensity
+    f_occ = torch.clamp(f_occ, 0.0, 1.0)
+    
+    return f_occ.unsqueeze(-1)
+
+
+def aces_tonemapping(x: torch.Tensor) -> torch.Tensor:
+    """
+    Applies ACES tone mapping curve to an HDR image tensor.
+    Input:  x - HDR tensor, shape (..., 3), range [0, +inf)
+    Output: LDR tensor, same shape, range [0, 1]
+    
+    NOTE: This causes bright base_color objects to over-expose (appear white)
+    compared to Blender's standard sRGB display transform. Use linear_to_srgb()
+    instead for renders_cond alignment. See pbr_color_diff_debug_guide.md #5.
+    """
+    a = 2.51
+    b = 0.03
+    c = 2.43
+    d = 0.59
+    e = 0.14
+    
+    # Apply the ACES fitted curve
+    mapped = (x * (a * x + b)) / (x * (c * x + d) + e)
+    
+    # Clamp to [0, 1] for display or saving
+    return torch.clamp(mapped, 0.0, 1.0)
+
+
+def gamma_correction(x: torch.Tensor, gamma: float = 2.2) -> torch.Tensor:
+    """
+    Applies gamma correction to an HDR image tensor.
+    """
+    return torch.clamp(x ** (1.0 / gamma), 0.0, 1.0)
+
+
+def linear_to_srgb(x: torch.Tensor) -> torch.Tensor:
+    """
+    Standard linear-to-sRGB conversion matching Blender's sRGB display transform.
+    
+    Applies the official sRGB EOTF inverse:
+      - For values <= 0.0031308: sRGB = 12.92 * linear
+      - For values >  0.0031308: sRGB = 1.055 * linear^(1/2.4) - 0.055
+    
+    Input:  x - linear HDR tensor, clamped to [0, +inf)
+    Output: sRGB tensor, range [0, 1]
+    """
+    x = torch.clamp(x, 0.0)
+    low = 12.92 * x
+    high = 1.055 * torch.pow(torch.clamp(x, min=0.0031308), 1.0 / 2.4) - 0.055
+    return torch.clamp(torch.where(x <= 0.0031308, low, high), 0.0, 1.0)
+    
+
+class PbrMeshRenderer:
+    """
+    Renderer for the PBR mesh.
+
+    Args:
+        rendering_options (dict): Rendering options.
+        """
+    def __init__(self, rendering_options={}, device='cuda'):
+        if 'dr' not in globals():
+            import nvdiffrast.torch as dr
+        
+        self.rendering_options = edict({
+            "resolution": None,
+            "near": None,
+            "far": None,
+            "ssaa": 1,
+            "peel_layers": 8,
+        })
+        self.rendering_options.update(rendering_options)
+        self.glctx = dr.RasterizeCudaContext(device=device)
+        self.device=device
+        
+    def render(
+            self,
+            mesh : Mesh,
+            extrinsics: torch.Tensor,
+            intrinsics: torch.Tensor,
+            envmap : Union[EnvMap, Dict[str, EnvMap]],
+            use_envmap_bg : bool = False,
+            transformation : Optional[torch.Tensor] = None
+        ) -> edict:
+        """
+        Render the mesh.
+
+        Args:
+            mesh : meshmodel
+            extrinsics (torch.Tensor): (4, 4) camera extrinsics
+            intrinsics (torch.Tensor): (3, 3) camera intrinsics
+            envmap (Union[EnvMap, Dict[str, EnvMap]]): environment map or a dictionary of environment maps
+            use_envmap_bg (bool): whether to use envmap as background
+            transformation (torch.Tensor): (4, 4) transformation matrix
+
+        Returns:
+            edict based on return_types containing:
+                shaded (torch.Tensor): [3, H, W] shaded color image
+                normal (torch.Tensor): [3, H, W] normal image
+                base_color (torch.Tensor): [3, H, W] base color image
+                metallic (torch.Tensor): [H, W] metallic image
+                roughness (torch.Tensor): [H, W] roughness image
+        """
+        if 'dr' not in globals():
+            import nvdiffrast.torch as dr
+            
+        if not isinstance(envmap, dict):
+            envmap = {'' : envmap}
+        num_envmaps = len(envmap)
+            
+        resolution = self.rendering_options["resolution"]
+        near = self.rendering_options["near"]
+        far = self.rendering_options["far"]
+        ssaa = self.rendering_options["ssaa"]
+        
+        if mesh.vertices.shape[0] == 0 or mesh.faces.shape[0] == 0 or \
+           not torch.isfinite(mesh.vertices).all() or \
+           mesh.faces.max() >= mesh.vertices.shape[0] or mesh.faces.min() < 0:
+            if mesh.vertices.shape[0] > 0 and not torch.isfinite(mesh.vertices).all():
+                print(f"[PbrMeshRenderer] WARNING: mesh vertices contain NaN/Inf, returning blank image.")
+            if mesh.faces.shape[0] > 0 and mesh.vertices.shape[0] > 0 and \
+               (mesh.faces.max() >= mesh.vertices.shape[0] or mesh.faces.min() < 0):
+                print(f"[PbrMeshRenderer] WARNING: mesh faces contain out-of-bound indices "
+                      f"(max={mesh.faces.max().item()}, min={mesh.faces.min().item()}, "
+                      f"num_vertices={mesh.vertices.shape[0]}), returning blank image.")
+            out_dict = edict(
+                normal=torch.zeros((3, resolution, resolution), dtype=torch.float32, device=self.device),
+                mask=torch.zeros((resolution, resolution), dtype=torch.float32, device=self.device),
+                base_color=torch.zeros((3, resolution, resolution), dtype=torch.float32, device=self.device),
+                metallic=torch.zeros((resolution, resolution), dtype=torch.float32, device=self.device),
+                roughness=torch.zeros((resolution, resolution), dtype=torch.float32, device=self.device),
+                alpha=torch.zeros((resolution, resolution), dtype=torch.float32, device=self.device),
+                clay=torch.zeros((resolution, resolution), dtype=torch.float32, device=self.device),
+            )
+            for i, k in enumerate(envmap.keys()):
+                shaded_key = f"shaded_{k}" if k != '' else "shaded"
+                out_dict[shaded_key] = torch.zeros((3, resolution, resolution), dtype=torch.float32, device=self.device)
+            return out_dict
+            
+        rays_o, rays_d = utils3d.torch.get_image_rays(
+            extrinsics, intrinsics, resolution * ssaa, resolution * ssaa
+        )
+        
+        perspective = intrinsics_to_projection(intrinsics, near, far)
+        
+        full_proj = (perspective @ extrinsics).unsqueeze(0)
+        extrinsics = extrinsics.unsqueeze(0)
+        
+        vertices = mesh.vertices.unsqueeze(0)
+        vertices_orig = vertices.clone()
+        vertices_homo = torch.cat([vertices, torch.ones_like(vertices[..., :1])], dim=-1)
+        if transformation is not None:
+            vertices_homo = torch.bmm(vertices_homo, transformation.unsqueeze(0).transpose(-1, -2))
+            vertices = vertices_homo[..., :3].contiguous()
+        vertices_camera = torch.bmm(vertices_homo, extrinsics.transpose(-1, -2))
+        vertices_clip = torch.bmm(vertices_homo, full_proj.transpose(-1, -2))
+        faces = mesh.faces
+        
+        v0 = vertices[0, mesh.faces[:, 0], :3]
+        v1 = vertices[0, mesh.faces[:, 1], :3]
+        v2 = vertices[0, mesh.faces[:, 2], :3]
+        e0 = v1 - v0
+        e1 = v2 - v0
+        face_normal = torch.cross(e0, e1, dim=1)
+        face_normal = F.normalize(face_normal, dim=1)
+        
+        out_dict = edict()
+        shaded = torch.zeros((num_envmaps, resolution * ssaa, resolution * ssaa, 3), dtype=torch.float32, device=self.device)
+        depth = torch.full((resolution * ssaa, resolution * ssaa, 1), 1e10, dtype=torch.float32, device=self.device)
+        normal = torch.zeros((resolution * ssaa, resolution * ssaa, 3), dtype=torch.float32, device=self.device)
+        max_w = torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device)
+        alpha = torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device)
+        with dr.DepthPeeler(self.glctx, vertices_clip, faces, (resolution * ssaa, resolution * ssaa)) as peeler:
+            for _ in range(self.rendering_options["peel_layers"]):
+                rast, rast_db = peeler.rasterize_next_layer()
+                
+                # Pos
+                pos = dr.interpolate(vertices, rast, faces)[0][0]
+                
+                # Depth
+                gb_depth = dr.interpolate(vertices_camera[..., 2:3].contiguous(), rast, faces)[0][0]
+                        
+                # Normal
+                gb_normal = dr.interpolate(face_normal.unsqueeze(0), rast, torch.arange(face_normal.shape[0], dtype=torch.int, device=self.device).unsqueeze(1).repeat(1, 3).contiguous())[0][0]
+                gb_normal = torch.where(
+                    torch.sum(gb_normal * (pos - rays_o), dim=-1, keepdim=True) > 0,
+                    -gb_normal,
+                    gb_normal
+                )
+                gb_cam_normal = (extrinsics[..., :3, :3].reshape(1, 1, 3, 3) @ gb_normal.unsqueeze(-1)).squeeze(-1)
+                if _ == 0:
+                    out_dict.normal = -gb_cam_normal * 0.5 + 0.5
+                    mask = (rast[0, ..., -1:] > 0).float()
+                    out_dict.mask = mask
+                
+                # PBR attributes
+                if isinstance(mesh, MeshWithVoxel):
+                    if 'grid_sample_3d' not in globals():
+                        from flex_gemm.ops.grid_sample import grid_sample_3d
+                    mask = rast[..., -1:] > 0
+                    xyz = dr.interpolate(vertices_orig, rast, faces)[0]
+                    xyz = ((xyz - mesh.origin) / mesh.voxel_size).reshape(1, -1, 3)
+                    img = grid_sample_3d(
+                        mesh.attrs,
+                        torch.cat([torch.zeros_like(mesh.coords[..., :1]), mesh.coords], dim=-1),
+                        mesh.voxel_shape,
+                        xyz,
+                        mode='trilinear'
+                    )
+                    img = img.reshape(1, resolution * ssaa, resolution * ssaa, mesh.attrs.shape[-1]) * mask
+                    gb_basecolor = img[0, ..., mesh.layout['base_color']]
+                    gb_metallic = img[0, ..., mesh.layout['metallic']]
+                    gb_roughness = img[0, ..., mesh.layout['roughness']]
+                    gb_alpha = img[0, ..., mesh.layout['alpha']]
+                elif isinstance(mesh, MeshWithPbrMaterial):
+                    tri_id = rast[0, :, :, -1:]
+                    mask = tri_id > 0
+                    uv_coords = mesh.uv_coords.reshape(1, -1, 2)
+                    texc, texd = dr.interpolate(
+                        uv_coords,
+                        rast,
+                        torch.arange(mesh.uv_coords.shape[0] * 3, dtype=torch.int, device=self.device).reshape(-1, 3),
+                        rast_db=rast_db,
+                        diff_attrs='all'
+                    )
+                    # Fix problematic texture coordinates
+                    texc = torch.nan_to_num(texc, nan=0.0, posinf=1e3, neginf=-1e3)
+                    texc = torch.clamp(texc, min=-1e3, max=1e3)
+                    texd = torch.nan_to_num(texd, nan=0.0, posinf=1e3, neginf=-1e3)
+                    texd = torch.clamp(texd, min=-1e3, max=1e3)
+                    mid = mesh.material_ids[(tri_id - 1).long()]
+                    gb_basecolor = torch.zeros((resolution * ssaa, resolution * ssaa, 3), dtype=torch.float32, device=self.device)
+                    gb_metallic = torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device)
+                    gb_roughness = torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device)
+                    gb_alpha = torch.zeros((resolution * ssaa, resolution * ssaa, 1), dtype=torch.float32, device=self.device)
+                    for id, mat in enumerate(mesh.materials):
+                        mat_mask = (mid == id).float() * mask.float()
+                        mat_texc = texc * mat_mask
+                        mat_texd = texd * mat_mask
+
+                        if mat.base_color_texture is not None:
+                            bc = dr.texture(
+                                mat.base_color_texture.image.unsqueeze(0),
+                                mat_texc,
+                                mat_texd,
+                                filter_mode='linear-mipmap-linear' if mat.base_color_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest',
+                                boundary_mode='clamp' if mat.base_color_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap'
+                            )[0]
+                            gb_basecolor += bc * mat.base_color_factor * mat_mask
+                        else:
+                            gb_basecolor += mat.base_color_factor * mat_mask
+                            
+                        if mat.metallic_texture is not None:
+                            m = dr.texture(
+                                mat.metallic_texture.image.unsqueeze(0),
+                                mat_texc,
+                                mat_texd,
+                                filter_mode='linear-mipmap-linear' if mat.metallic_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest',
+                                boundary_mode='clamp' if mat.metallic_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap'
+                            )[0]
+                            gb_metallic += m * mat.metallic_factor * mat_mask
+                        else:
+                            gb_metallic += mat.metallic_factor * mat_mask
+
+                        if mat.roughness_texture is not None:
+                            r = dr.texture(
+                                mat.roughness_texture.image.unsqueeze(0),
+                                mat_texc,
+                                mat_texd,
+                                filter_mode='linear-mipmap-linear' if mat.roughness_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest',
+                                boundary_mode='clamp' if mat.roughness_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap'
+                            )[0]
+                            gb_roughness += r * mat.roughness_factor * mat_mask
+                        else:
+                            gb_roughness += mat.roughness_factor * mat_mask
+
+                        if mat.alpha_mode == AlphaMode.OPAQUE:
+                            gb_alpha += 1.0 * mat_mask
+                        else:
+                            if mat.alpha_texture is not None:
+                                a = dr.texture(
+                                    mat.alpha_texture.image.unsqueeze(0),
+                                    mat_texc,
+                                    mat_texd,
+                                    filter_mode='linear-mipmap-linear' if mat.alpha_texture.filter_mode == TextureFilterMode.LINEAR else 'nearest',
+                                    boundary_mode='clamp' if mat.alpha_texture.wrap_mode == TextureWrapMode.CLAMP_TO_EDGE else 'wrap'
+                                )[0]
+                                if mat.alpha_mode == AlphaMode.MASK:
+                                    gb_alpha += (a * mat.alpha_factor > mat.alpha_cutoff).float() * mat_mask
+                                elif mat.alpha_mode == AlphaMode.BLEND:
+                                    gb_alpha += a * mat.alpha_factor * mat_mask
+                            else:
+                                if mat.alpha_mode == AlphaMode.MASK:
+                                    gb_alpha += (mat.alpha_factor > mat.alpha_cutoff).float() * mat_mask
+                                elif mat.alpha_mode == AlphaMode.BLEND:
+                                    gb_alpha += mat.alpha_factor * mat_mask
+                if _ == 0:
+                    out_dict.base_color = gb_basecolor
+                    out_dict.metallic = gb_metallic
+                    out_dict.roughness = gb_roughness
+                    out_dict.alpha = gb_alpha
+                    
+                # Shading
+                #TODO
+                gb_basecolor = torch.clamp(gb_basecolor, 0.0, 1.0) ** 2.2
+                # gb_basecolor = torch.clamp(gb_basecolor, 0.0, 1.0)
+                gb_metallic = torch.clamp(gb_metallic, 0.0, 1.0)
+                gb_roughness = torch.clamp(gb_roughness, 0.0, 1.0)
+                gb_alpha = torch.clamp(gb_alpha, 0.0, 1.0)
+                gb_orm = torch.cat([
+                    torch.zeros_like(gb_metallic),
+                    gb_roughness,
+                    gb_metallic,
+                ], dim=-1)
+                gb_shaded = torch.stack([
+                    e.shade(
+                        pos.unsqueeze(0),
+                        gb_normal.unsqueeze(0),
+                        gb_basecolor.unsqueeze(0),
+                        gb_orm.unsqueeze(0),
+                        rays_o,
+                        specular=True,
+                    )[0]
+                    for e in envmap.values()
+                ], dim=0)
+                
+                # Compositing
+                w = (1 - alpha) * gb_alpha
+                depth = torch.where(w > max_w, gb_depth, depth)
+                normal = torch.where(w > max_w, gb_cam_normal, normal)
+                max_w = torch.maximum(max_w, w)
+                shaded += w * gb_shaded
+                alpha += w
+        
+        # Ambient occulusion
+        f_occ = screen_space_ambient_occlusion(
+            depth, normal, perspective, intensity=1.5
+        )
+        shaded *= (1 - f_occ)
+        out_dict.clay = (1 - f_occ)
+                
+        # Background
+        if use_envmap_bg:
+            bg = torch.stack([e.sample(rays_d) for e in envmap.values()], dim=0)
+            shaded += (1 - alpha) * bg
+        
+        for i, k in enumerate(envmap.keys()):
+            shaded_key = f"shaded_{k}" if k != '' else "shaded"
+            out_dict[shaded_key] = shaded[i]
+    
+        # SSAA
+        for k in out_dict.keys():
+            if ssaa > 1:
+                out_dict[k] = F.interpolate(out_dict[k].unsqueeze(0).permute(0, 3, 1, 2), (resolution, resolution), mode='bilinear', align_corners=False, antialias=True)
+            else:
+                out_dict[k] = out_dict[k].permute(2, 0, 1)
+            out_dict[k] = out_dict[k].squeeze()
+                
+        # Post processing: linear → sRGB (matches Blender's display transform)
+        for k in envmap.keys():
+            shaded_key = f"shaded_{k}" if k != '' else "shaded"
+            out_dict[shaded_key] = linear_to_srgb(out_dict[shaded_key])
+            
+        return out_dict
diff --git a/pixal3d/renderers/voxel_renderer.py b/pixal3d/renderers/voxel_renderer.py
new file mode 100644
index 0000000..dfe28ad
--- /dev/null
+++ b/pixal3d/renderers/voxel_renderer.py
@@ -0,0 +1,68 @@
+import torch
+from easydict import EasyDict as edict
+from ..representations import Voxel
+from easydict import EasyDict as edict
+
+
+class VoxelRenderer:
+    """
+    Renderer for the Voxel representation.
+
+    Args:
+        rendering_options (dict): Rendering options.
+    """
+
+    def __init__(self, rendering_options={}) -> None:
+        self.rendering_options = edict({
+            "resolution": None,
+            "near": 0.1,
+            "far": 10.0,
+            "ssaa": 1,
+        })
+        self.rendering_options.update(rendering_options)
+    
+    def render(
+            self,
+            voxel: Voxel,
+            extrinsics: torch.Tensor,
+            intrinsics: torch.Tensor,
+            colors_overwrite: torch.Tensor = None
+        ) -> edict:
+        """
+        Render the gausssian.
+
+        Args:
+            voxel (Voxel): Voxel representation.
+            extrinsics (torch.Tensor): (4, 4) camera extrinsics
+            intrinsics (torch.Tensor): (3, 3) camera intrinsics
+            colors_overwrite (torch.Tensor): (N, 3) override color
+
+        Returns:
+            edict containing:
+                color (torch.Tensor): (3, H, W) rendered color image
+                depth (torch.Tensor): (H, W) rendered depth
+                alpha (torch.Tensor): (H, W) rendered alpha
+                ...
+        """ 
+        # lazy import
+        if 'o_voxel' not in globals():
+            import o_voxel
+        renderer = o_voxel.rasterize.VoxelRenderer(self.rendering_options)
+        positions = voxel.position
+        attrs = voxel.attrs if colors_overwrite is None else colors_overwrite
+        voxel_size = voxel.voxel_size
+        
+        # Render
+        render_ret = renderer.render(positions, attrs, voxel_size, extrinsics, intrinsics)
+        
+        ret = {
+            'depth': render_ret['depth'],
+            'alpha': render_ret['alpha'],
+        }
+        if colors_overwrite is not None:
+            ret['color'] = render_ret['attr']
+        else:
+            for k, s in voxel.layout.items():
+                ret[k] = render_ret['attr'][s]
+        
+        return ret
diff --git a/pixal3d/representations/__init__.py b/pixal3d/representations/__init__.py
new file mode 100644
index 0000000..0e7d929
--- /dev/null
+++ b/pixal3d/representations/__init__.py
@@ -0,0 +1,31 @@
+import importlib
+
+__attributes = {
+    'Mesh': 'mesh',
+    'Voxel': 'voxel',
+    'MeshWithVoxel': 'mesh',
+    'MeshWithPbrMaterial': 'mesh',
+}
+
+__submodules = []
+
+__all__ = list(__attributes.keys()) + __submodules
+
+def __getattr__(name):
+    if name not in globals():
+        if name in __attributes:
+            module_name = __attributes[name]
+            module = importlib.import_module(f".{module_name}", __name__)
+            globals()[name] = getattr(module, name)
+        elif name in __submodules:
+            module = importlib.import_module(f".{name}", __name__)
+            globals()[name] = module
+        else:
+            raise AttributeError(f"module {__name__} has no attribute {name}")
+    return globals()[name]
+
+
+# For Pylance
+if __name__ == '__main__':
+    from .mesh import Mesh, MeshWithVoxel, MeshWithPbrMaterial
+    from .voxel import Voxel
diff --git a/pixal3d/representations/mesh/__init__.py b/pixal3d/representations/mesh/__init__.py
new file mode 100644
index 0000000..aff4c99
--- /dev/null
+++ b/pixal3d/representations/mesh/__init__.py
@@ -0,0 +1 @@
+from .base import Mesh, MeshWithVoxel, MeshWithPbrMaterial, TextureFilterMode, TextureWrapMode, AlphaMode, PbrMaterial, Texture
diff --git a/pixal3d/representations/mesh/base.py b/pixal3d/representations/mesh/base.py
new file mode 100644
index 0000000..6bdb47f
--- /dev/null
+++ b/pixal3d/representations/mesh/base.py
@@ -0,0 +1,234 @@
+from typing import *
+import torch
+from ..voxel import Voxel
+import cumesh
+from flex_gemm.ops.grid_sample import grid_sample_3d
+
+
+class Mesh:
+    def __init__(self,
+        vertices,
+        faces,
+        vertex_attrs=None
+    ):
+        self.vertices = vertices.float()
+        self.faces = faces.int()
+        self.vertex_attrs = vertex_attrs
+        
+    @property
+    def device(self):
+        return self.vertices.device
+        
+    def to(self, device, non_blocking=False):
+        return Mesh(
+            self.vertices.to(device, non_blocking=non_blocking),
+            self.faces.to(device, non_blocking=non_blocking),
+            self.vertex_attrs.to(device, non_blocking=non_blocking) if self.vertex_attrs is not None else None,
+        )
+        
+    def cuda(self, non_blocking=False):
+        return self.to('cuda', non_blocking=non_blocking)
+        
+    def cpu(self):
+        return self.to('cpu')
+    
+    def fill_holes(self, max_hole_perimeter=3e-2):
+        vertices = self.vertices.clone().cuda().contiguous()
+        faces = self.faces.clone().cuda().contiguous()
+        
+        mesh = cumesh.CuMesh()
+        mesh.init(vertices, faces)
+        mesh.get_edges()
+        mesh.get_boundary_info()
+        if mesh.num_boundaries == 0:
+            return
+        mesh.get_vertex_edge_adjacency()
+        mesh.get_vertex_boundary_adjacency()
+        mesh.get_manifold_boundary_adjacency()
+        mesh.read_manifold_boundary_adjacency()
+        mesh.get_boundary_connected_components()
+        mesh.get_boundary_loops()
+        if mesh.num_boundary_loops == 0:
+            return
+        mesh.fill_holes(max_hole_perimeter=max_hole_perimeter)
+        new_vertices, new_faces = mesh.read()
+        
+        self.vertices = new_vertices.to(self.device)
+        self.faces = new_faces.to(self.device)
+        
+    def remove_faces(self, face_mask: torch.Tensor):
+        vertices = self.vertices.clone().cuda().contiguous()
+        faces = self.faces.clone().cuda().contiguous()
+        
+        mesh = cumesh.CuMesh()
+        mesh.init(vertices, faces)
+        mesh.remove_faces(face_mask)
+        new_vertices, new_faces = mesh.read()
+        
+        self.vertices = new_vertices.to(self.device)
+        self.faces = new_faces.to(self.device)
+        
+    def simplify(self, target=1000000, verbose: bool=False, options: dict={}):
+        vertices = self.vertices.clone().cuda().contiguous()
+        faces = self.faces.clone().cuda().contiguous()
+        
+        mesh = cumesh.CuMesh()
+        mesh.init(vertices, faces)
+        mesh.simplify(target, verbose=verbose, options=options)
+        new_vertices, new_faces = mesh.read()
+        
+        self.vertices = new_vertices.to(self.device)
+        self.faces = new_faces.to(self.device)
+
+
+class TextureFilterMode:
+    CLOSEST = 0
+    LINEAR = 1
+
+
+class TextureWrapMode:
+    CLAMP_TO_EDGE = 0
+    REPEAT = 1
+    MIRRORED_REPEAT = 2
+
+
+class AlphaMode:
+    OPAQUE = 0
+    MASK = 1
+    BLEND = 2
+
+
+class Texture:
+    def __init__(
+        self,
+        image: torch.Tensor,
+        filter_mode: TextureFilterMode = TextureFilterMode.LINEAR,
+        wrap_mode: TextureWrapMode = TextureWrapMode.REPEAT
+    ):
+        self.image = image
+        self.filter_mode = filter_mode
+        self.wrap_mode = wrap_mode
+
+    def to(self, device, non_blocking=False):
+        return Texture(
+            self.image.to(device, non_blocking=non_blocking),
+            self.filter_mode,
+            self.wrap_mode,
+        )
+
+
+class PbrMaterial:
+    def __init__(
+        self,
+        base_color_texture: Optional[Texture] = None,
+        base_color_factor: Union[torch.Tensor, List[float]] = [1.0, 1.0, 1.0],
+        metallic_texture: Optional[Texture] = None,
+        metallic_factor: float = 1.0,
+        roughness_texture: Optional[Texture] = None,
+        roughness_factor: float = 1.0,
+        alpha_texture: Optional[Texture] = None,
+        alpha_factor: float = 1.0,
+        alpha_mode: AlphaMode = AlphaMode.OPAQUE,
+        alpha_cutoff: float = 0.5,
+    ):
+        self.base_color_texture = base_color_texture
+        self.base_color_factor = torch.tensor(base_color_factor, dtype=torch.float32)[:3]
+        self.metallic_texture = metallic_texture
+        self.metallic_factor = metallic_factor
+        self.roughness_texture = roughness_texture
+        self.roughness_factor = roughness_factor
+        self.alpha_texture = alpha_texture
+        self.alpha_factor = alpha_factor
+        self.alpha_mode = alpha_mode
+        self.alpha_cutoff = alpha_cutoff
+
+    def to(self, device, non_blocking=False):
+        return PbrMaterial(
+            base_color_texture=self.base_color_texture.to(device, non_blocking=non_blocking) if self.base_color_texture is not None else None,
+            base_color_factor=self.base_color_factor.to(device, non_blocking=non_blocking),
+            metallic_texture=self.metallic_texture.to(device, non_blocking=non_blocking) if self.metallic_texture is not None else None,
+            metallic_factor=self.metallic_factor,
+            roughness_texture=self.roughness_texture.to(device, non_blocking=non_blocking) if self.roughness_texture is not None else None,
+            roughness_factor=self.roughness_factor,
+            alpha_texture=self.alpha_texture.to(device, non_blocking=non_blocking) if self.alpha_texture is not None else None,
+            alpha_factor=self.alpha_factor,
+            alpha_mode=self.alpha_mode,
+            alpha_cutoff=self.alpha_cutoff,
+        )
+
+
+class MeshWithPbrMaterial(Mesh):
+    def __init__(self,
+        vertices,
+        faces,
+        material_ids,
+        uv_coords,
+        materials: List[PbrMaterial],
+    ):
+        self.vertices = vertices.float()
+        self.faces = faces.int()
+        self.material_ids = material_ids    # [M]
+        self.uv_coords = uv_coords          # [M, 3, 2]
+        self.materials = materials
+        self.layout = {
+            'base_color': slice(0, 3),
+            'metallic': slice(3, 4),
+            'roughness': slice(4, 5),
+            'alpha': slice(5, 6),
+        }
+
+    def to(self, device, non_blocking=False):
+        return MeshWithPbrMaterial(
+            self.vertices.to(device, non_blocking=non_blocking),
+            self.faces.to(device, non_blocking=non_blocking),
+            self.material_ids.to(device, non_blocking=non_blocking),
+            self.uv_coords.to(device, non_blocking=non_blocking),
+            [material.to(device, non_blocking=non_blocking) for material in self.materials],
+        )
+
+
+class MeshWithVoxel(Mesh, Voxel):
+    def __init__(self,
+        vertices: torch.Tensor,
+        faces: torch.Tensor,
+        origin: list,
+        voxel_size: float,
+        coords: torch.Tensor,
+        attrs: torch.Tensor,
+        voxel_shape: torch.Size,
+        layout: Dict = {},
+    ):
+        self.vertices = vertices.float()
+        self.faces = faces.int()
+        self.origin = torch.tensor(origin, dtype=torch.float32, device=self.device)
+        self.voxel_size = voxel_size
+        self.coords = coords
+        self.attrs = attrs
+        self.voxel_shape = voxel_shape
+        self.layout = layout
+
+    def to(self, device, non_blocking=False):
+        return MeshWithVoxel(
+            self.vertices.to(device, non_blocking=non_blocking),
+            self.faces.to(device, non_blocking=non_blocking),
+            self.origin.tolist(),
+            self.voxel_size,
+            self.coords.to(device, non_blocking=non_blocking),
+            self.attrs.to(device, non_blocking=non_blocking),
+            self.voxel_shape,
+            self.layout,
+        )
+        
+    def query_attrs(self, xyz):
+        grid = ((xyz - self.origin) / self.voxel_size).reshape(1, -1, 3)
+        vertex_attrs = grid_sample_3d(
+            self.attrs,
+            torch.cat([torch.zeros_like(self.coords[..., :1]), self.coords], dim=-1),
+            self.voxel_shape,
+            grid,
+            mode='trilinear'
+        )[0]
+        return vertex_attrs
+        
+    def query_vertex_attrs(self):
+        return self.query_attrs(self.vertices)
diff --git a/pixal3d/representations/voxel/__init__.py b/pixal3d/representations/voxel/__init__.py
new file mode 100644
index 0000000..b5792ea
--- /dev/null
+++ b/pixal3d/representations/voxel/__init__.py
@@ -0,0 +1 @@
+from .voxel_model import Voxel
\ No newline at end of file
diff --git a/pixal3d/representations/voxel/voxel_model.py b/pixal3d/representations/voxel/voxel_model.py
new file mode 100644
index 0000000..9317ab2
--- /dev/null
+++ b/pixal3d/representations/voxel/voxel_model.py
@@ -0,0 +1,54 @@
+from typing import Dict
+import torch
+
+
+class Voxel:
+    def __init__(
+            self, 
+            origin: list,
+            voxel_size: float,
+            coords: torch.Tensor = None,
+            attrs: torch.Tensor = None,
+            layout: Dict = {},
+            device: torch.device = 'cuda'
+        ):
+        self.origin = torch.tensor(origin, dtype=torch.float32, device=device)
+        self.voxel_size = voxel_size
+        self.coords = coords
+        self.attrs = attrs
+        self.layout = layout
+        self.device = device
+        
+    @property
+    def position(self):
+        return (self.coords + 0.5) * self.voxel_size + self.origin[None, :]
+    
+    def split_attrs(self):
+        return {
+            k: self.attrs[:, self.layout[k]]
+            for k in self.layout
+        }
+        
+    def save(self, path):
+        # lazy import
+        if 'o_voxel' not in globals():
+            import o_voxel
+        o_voxel.io.write(
+            path,
+            self.coords,
+            self.split_attrs(),
+        )
+        
+    def load(self, path):
+        # lazy import
+        if 'o_voxel' not in globals():
+            import o_voxel
+        coord, attrs = o_voxel.io.read(path)
+        self.coords = coord.int().to(self.device)
+        self.attrs = torch.cat([attrs[k] for k in attrs], dim=1).to(self.device)
+        # build layout
+        start = 0
+        self.layout = {}
+        for k in attrs:
+            self.layout[k] = slice(start, start + attrs[k].shape[1])
+            start += attrs[k].shape[1]
diff --git a/pixal3d/trainers/__init__.py b/pixal3d/trainers/__init__.py
new file mode 100644
index 0000000..7a10751
--- /dev/null
+++ b/pixal3d/trainers/__init__.py
@@ -0,0 +1,74 @@
+import importlib
+
+__attributes = {
+    'BasicTrainer': 'basic',
+    
+    'SparseStructureVaeTrainer': 'vae.sparse_structure_vae',
+    'ShapeVaeTrainer': 'vae.shape_vae',
+    'PbrVaeTrainer': 'vae.pbr_vae',
+    
+    'FlowMatchingTrainer': 'flow_matching.flow_matching',
+    'FlowMatchingCFGTrainer': 'flow_matching.flow_matching',
+    'TextConditionedFlowMatchingCFGTrainer': 'flow_matching.flow_matching',
+    'ImageConditionedFlowMatchingCFGTrainer': 'flow_matching.flow_matching',
+    'ImageConditionedProjFlowMatchingCFGTrainer': 'flow_matching.flow_matching',
+    
+    'SparseFlowMatchingTrainer': 'flow_matching.sparse_flow_matching',
+    'SparseFlowMatchingCFGTrainer': 'flow_matching.sparse_flow_matching',
+    'TextConditionedSparseFlowMatchingCFGTrainer': 'flow_matching.sparse_flow_matching',
+    'ImageConditionedSparseFlowMatchingCFGTrainer': 'flow_matching.sparse_flow_matching',
+    'MultiImageConditionedSparseFlowMatchingCFGTrainer': 'flow_matching.sparse_flow_matching',
+    'ImageConditionedProjSparseFlowMatchingCFGTrainer': 'flow_matching.sparse_flow_matching',
+    
+    'DinoV2FeatureExtractor': 'flow_matching.mixins.image_conditioned',
+    'DinoV3FeatureExtractor': 'flow_matching.mixins.image_conditioned',
+    'DinoV3ProjFeatureExtractor': 'flow_matching.mixins.image_conditioned_proj',
+    'DinoV3VaeProjFeatureExtractor': 'flow_matching.mixins.image_conditioned_proj',
+    'ImageConditionedProjMixin': 'flow_matching.mixins.image_conditioned_proj',
+}
+
+__submodules = []
+
+__all__ = list(__attributes.keys()) + __submodules
+
+def __getattr__(name):
+    if name not in globals():
+        if name in __attributes:
+            module_name = __attributes[name]
+            module = importlib.import_module(f".{module_name}", __name__)
+            globals()[name] = getattr(module, name)
+        elif name in __submodules:
+            module = importlib.import_module(f".{name}", __name__)
+            globals()[name] = module
+        else:
+            raise AttributeError(f"module {__name__} has no attribute {name}")
+    return globals()[name]
+
+
+# For Pylance
+if __name__ == '__main__':
+    from .basic import BasicTrainer
+
+    from .vae.sparse_structure_vae import SparseStructureVaeTrainer
+    from .vae.shape_vae import ShapeVaeTrainer
+    from .vae.pbr_vae import PbrVaeTrainer
+    
+    from .flow_matching.flow_matching import (
+        FlowMatchingTrainer,
+        FlowMatchingCFGTrainer,
+        TextConditionedFlowMatchingCFGTrainer,
+        ImageConditionedFlowMatchingCFGTrainer,
+        ImageConditionedProjFlowMatchingCFGTrainer,
+    )
+    
+    from .flow_matching.sparse_flow_matching import (
+        SparseFlowMatchingTrainer,
+        SparseFlowMatchingCFGTrainer,
+        TextConditionedSparseFlowMatchingCFGTrainer,
+        ImageConditionedSparseFlowMatchingCFGTrainer,
+    )
+    
+    from .flow_matching.mixins.image_conditioned import (
+        DinoV2FeatureExtractor,
+        DinoV3FeatureExtractor,
+    )
diff --git a/pixal3d/trainers/basic.py b/pixal3d/trainers/basic.py
new file mode 100644
index 0000000..2846b8f
--- /dev/null
+++ b/pixal3d/trainers/basic.py
@@ -0,0 +1,1293 @@
+from abc import abstractmethod
+import os
+import time
+import json
+import copy
+import threading
+from functools import partial
+from contextlib import nullcontext
+
+import torch
+import torch.distributed as dist
+from torch.utils.data import DataLoader
+from torch.nn.parallel import DistributedDataParallel as DDP
+import numpy as np
+
+from torchvision import utils
+
+try:
+    import wandb
+    WANDB_AVAILABLE = True
+except ImportError:
+    WANDB_AVAILABLE = False
+
+from .utils import *
+from ..utils.general_utils import *
+from ..utils.data_utils import recursive_to_device, cycle, ResumableSampler
+from ..utils.dist_utils import *
+from ..utils import grad_clip_utils, elastic_utils
+
+
+class BasicTrainer:
+    """
+    Trainer for basic training loop.
+    
+    Args:
+        models (dict[str, nn.Module]): Models to train.
+        dataset (torch.utils.data.Dataset): Dataset.
+        output_dir (str): Output directory.
+        load_dir (str): Load directory.
+        step (int): Step to load.
+        batch_size (int): Batch size.
+        batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
+        batch_split (int): Split batch with gradient accumulation.
+        max_steps (int): Max steps.
+        optimizer (dict): Optimizer config.
+        lr_scheduler (dict): Learning rate scheduler config.
+        elastic (dict): Elastic memory management config.
+        grad_clip (float or dict): Gradient clip config.
+        ema_rate (float or list): Exponential moving average rates.
+        mix_precision_mode (str):
+            - None: No mixed precision.
+            - 'inflat_all': Hold a inflated fp32 master param for all params.
+            - 'amp': Automatic mixed precision.
+        mix_precision_dtype (str): Mixed precision dtype.
+        fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
+        parallel_mode (str): Parallel mode. Options are 'ddp'.
+        finetune_ckpt (dict): Finetune checkpoint.
+        log_param_stats (bool): Log parameter stats.
+        i_print (int): Print interval.
+        i_log (int): Log interval.
+        i_sample (int): Sample interval.
+        i_save (int): Save interval.
+        i_ddpcheck (int): DDP check interval.
+    """
+    def __init__(self,
+        models,
+        dataset,
+        *,
+        output_dir,
+        load_dir,
+        step,
+        max_steps,
+        batch_size=None,
+        batch_size_per_gpu=None,
+        batch_split=None,
+        optimizer={},
+        lr_scheduler=None,
+        elastic=None,
+        grad_clip=None,
+        ema_rate=0.9999,
+        fp16_mode=None,
+        mix_precision_mode='inflat_all',
+        mix_precision_dtype='float16',
+        fp16_scale_growth=1e-3,
+        parallel_mode='ddp',
+        finetune_ckpt=None,
+        log_param_stats=False,
+        prefetch_data=True,
+        snapshot_batch_size=4,
+        snapshot_num_samples=64,
+        num_workers=None,
+        debug=False,
+        i_print=1000,
+        i_log=500,
+        i_sample=10000,
+        i_save=10000,
+        i_ddpcheck=10000,
+        wandb_run=None,  # wandb run object
+        **kwargs
+    ):
+        assert batch_size is not None or batch_size_per_gpu is not None, 'Either batch_size or batch_size_per_gpu must be specified.'
+
+        self.models = models
+        self.dataset = dataset
+        self.batch_split = batch_split if batch_split is not None else 1
+        self.max_steps = max_steps
+        self.debug = debug
+        self.optimizer_config = optimizer
+        self.lr_scheduler_config = lr_scheduler
+        self.elastic_controller_config = elastic
+        self.grad_clip = grad_clip
+        self.ema_rate = [ema_rate] if isinstance(ema_rate, float) else ema_rate
+        if fp16_mode is not None:
+            mix_precision_dtype = 'float16'
+            mix_precision_mode = fp16_mode
+        self.mix_precision_mode = mix_precision_mode
+        self.mix_precision_dtype = str_to_dtype(mix_precision_dtype)
+        self.fp16_scale_growth = fp16_scale_growth
+        self.parallel_mode = parallel_mode
+        self.log_param_stats = log_param_stats
+        self.prefetch_data = prefetch_data
+        self.snapshot_batch_size = snapshot_batch_size
+        self.snapshot_num_samples = snapshot_num_samples
+        self.num_workers = num_workers
+        self.log = []
+        if self.prefetch_data:
+            self._data_prefetched = None
+
+        self.output_dir = output_dir
+        from datetime import datetime
+        self._log_file = os.path.join(self.output_dir, f'log_{datetime.now().strftime("%Y%m%d_%H%M%S")}.txt')
+        self.i_print = i_print
+        self.i_log = i_log
+        self.i_sample = i_sample
+        self.i_save = i_save
+        self.i_ddpcheck = i_ddpcheck        
+
+        if dist.is_initialized():
+            # Multi-GPU params
+            self.world_size = dist.get_world_size()
+            self.rank = dist.get_rank()
+            self.local_rank = dist.get_rank() % torch.cuda.device_count()
+            self.is_master = self.rank == 0
+        else:
+            # Single-GPU params
+            self.world_size = 1
+            self.rank = 0
+            self.local_rank = 0
+            self.is_master = True
+
+        self.batch_size = batch_size if batch_size_per_gpu is None else batch_size_per_gpu * self.world_size
+        self.batch_size_per_gpu = batch_size_per_gpu if batch_size_per_gpu is not None else batch_size // self.world_size
+        assert self.batch_size % self.world_size == 0, 'Batch size must be divisible by the number of GPUs.'
+        assert self.batch_size_per_gpu % self.batch_split == 0, 'Batch size per GPU must be divisible by batch split.'
+
+        self.init_models_and_more(**kwargs)
+        self.prepare_dataloader(**kwargs)
+        
+        # Load checkpoint
+        self.step = 0
+        if load_dir is not None and step is not None:
+            self.load(load_dir, step)
+        elif finetune_ckpt is not None:
+            self.finetune_from(finetune_ckpt)
+        
+        if self.is_master:
+            os.makedirs(os.path.join(self.output_dir, 'ckpts'), exist_ok=True)
+            os.makedirs(os.path.join(self.output_dir, 'samples'), exist_ok=True)
+            self.writer = None  # TensorBoard disabled (S3 FUSE does not support append)
+            # Initialize wandb
+            self.wandb_run = wandb_run
+            if self.wandb_run is not None:
+                print(f'Wandb logging enabled: {self.wandb_run.url}')
+
+        if self.parallel_mode == 'ddp' and self.world_size > 1:
+            self.check_ddp()
+            
+        if self.is_master:
+            print('\n\nTrainer initialized.')
+            print(self)
+
+    def __str__(self):
+        lines = []
+        lines.append(self.__class__.__name__)
+        lines.append(f'  - Models:')
+        for name, model in self.models.items():
+            lines.append(f'    - {name}: {model.__class__.__name__}')
+        lines.append(f'  - Dataset: {indent(str(self.dataset), 2)}')
+        lines.append(f'  - Dataloader:')
+        lines.append(f'    - Sampler: {self.dataloader.sampler.__class__.__name__}')
+        lines.append(f'    - Num workers: {self.dataloader.num_workers}')
+        lines.append(f'  - Number of steps: {self.max_steps}')
+        lines.append(f'  - Number of GPUs: {self.world_size}')
+        lines.append(f'  - Batch size: {self.batch_size}')
+        lines.append(f'  - Batch size per GPU: {self.batch_size_per_gpu}')
+        lines.append(f'  - Batch split: {self.batch_split}')
+        lines.append(f'  - Optimizer: {self.optimizer.__class__.__name__}')
+        lines.append(f'  - Learning rate: {self.optimizer.param_groups[0]["lr"]}')
+        if self.lr_scheduler_config is not None:
+            lines.append(f'  - LR scheduler: {self.lr_scheduler.__class__.__name__}')
+        if self.elastic_controller_config is not None:
+            lines.append(f'  - Elastic memory: {indent(str(self.elastic_controller), 2)}')
+        if self.grad_clip is not None:
+            lines.append(f'  - Gradient clip: {indent(str(self.grad_clip), 2)}')
+        lines.append(f'  - EMA rate: {self.ema_rate}')
+        lines.append(f'  - Mixed precision dtype: {self.mix_precision_dtype}')
+        lines.append(f'  - Mixed precision mode: {self.mix_precision_mode}')
+        if self.mix_precision_mode == 'amp' and self.mix_precision_dtype == torch.float16:
+            lines.append(f'  - FP16 scale growth: {self.fp16_scale_growth}')
+        lines.append(f'  - Parallel mode: {self.parallel_mode}')
+        return '\n'.join(lines)
+
+    @property
+    def device(self):
+        for _, model in self.models.items():
+            if hasattr(model, 'device'):
+                return model.device
+        return next(list(self.models.values())[0].parameters()).device
+            
+    def init_models_and_more(self, **kwargs):
+        """
+        Initialize models and more.
+        """
+        if self.world_size > 1:
+            # Prepare distributed data parallel
+            self.training_models = {
+                name: DDP(
+                    model,
+                    device_ids=[self.local_rank],
+                    output_device=self.local_rank,
+                    bucket_cap_mb=128,
+                    find_unused_parameters=False
+                )
+                for name, model in self.models.items()
+            }
+        else:
+            self.training_models = self.models
+
+        # Build master params
+        self.model_params = sum(
+            [[p for p in model.parameters() if p.requires_grad] for model in self.models.values()]
+        , [])
+        if self.mix_precision_mode == 'amp':
+            self.master_params = self.model_params
+            if self.mix_precision_dtype == torch.float16:
+                self.scaler = torch.GradScaler()
+        elif self.mix_precision_mode == 'inflat_all':
+            self.master_params = make_master_params(self.model_params)
+            if self.mix_precision_dtype == torch.float16:
+                self.log_scale = 20.0
+        elif self.mix_precision_mode is None:
+            self.master_params = self.model_params
+        else:
+            raise NotImplementedError(f'Mix precision mode {self.mix_precision_mode} is not implemented.')
+
+        # Build EMA params
+        if self.is_master:
+            self.ema_params = [copy.deepcopy(self.master_params) for _ in self.ema_rate]
+
+        # Initialize optimizer
+        if hasattr(torch.optim, self.optimizer_config['name']):
+            self.optimizer = getattr(torch.optim, self.optimizer_config['name'])(self.master_params, **self.optimizer_config['args'])
+        else:
+            self.optimizer = globals()[self.optimizer_config['name']](self.master_params, **self.optimizer_config['args'])
+        
+        # Initalize learning rate scheduler
+        if self.lr_scheduler_config is not None:
+            if hasattr(torch.optim.lr_scheduler, self.lr_scheduler_config['name']):
+                self.lr_scheduler = getattr(torch.optim.lr_scheduler, self.lr_scheduler_config['name'])(self.optimizer, **self.lr_scheduler_config['args'])
+            else:
+                self.lr_scheduler = globals()[self.lr_scheduler_config['name']](self.optimizer, **self.lr_scheduler_config['args'])
+
+        # Initialize elastic memory controller
+        if self.elastic_controller_config is not None:
+            assert any([isinstance(model, (elastic_utils.ElasticModule, elastic_utils.ElasticModuleMixin)) for model in self.models.values()]), \
+                'No elastic module found in models, please inherit from ElasticModule or ElasticModuleMixin'
+            self.elastic_controller = getattr(elastic_utils, self.elastic_controller_config['name'])(**self.elastic_controller_config['args'])
+            for model in self.models.values():
+                if isinstance(model, (elastic_utils.ElasticModule, elastic_utils.ElasticModuleMixin)):
+                    model.register_memory_controller(self.elastic_controller)
+
+        # Initialize gradient clipper
+        if self.grad_clip is not None:
+            if isinstance(self.grad_clip, (float, int)):
+                self.grad_clip = float(self.grad_clip)
+            else:
+                self.grad_clip = getattr(grad_clip_utils, self.grad_clip['name'])(**self.grad_clip['args'])
+
+    def prepare_dataloader(self, **kwargs):
+        """
+        Prepare dataloader.
+        """
+        self.data_sampler = ResumableSampler(
+            self.dataset,
+            shuffle=True,
+        )
+        if self.num_workers is None or self.num_workers == -1:
+            num_workers = max(1, int(np.ceil((os.cpu_count() - 16) / torch.cuda.device_count())))
+        else:
+            num_workers = self.num_workers
+        
+        self.dataloader = DataLoader(
+            self.dataset,
+            batch_size=self.batch_size_per_gpu,
+            num_workers=num_workers,
+            pin_memory=True,
+            drop_last=True,
+            persistent_workers=True,
+            collate_fn=self.dataset.collate_fn if hasattr(self.dataset, 'collate_fn') else None,
+            sampler=self.data_sampler,
+        )
+        self.data_iterator = cycle(self.dataloader)
+
+    def _master_params_to_state_dicts(self, master_params):
+        """
+        Convert master params to dict of state_dicts.
+        """
+        if self.mix_precision_mode == 'inflat_all':
+            master_params = unflatten_master_params(self.model_params, master_params)
+        state_dicts = {name: model.state_dict() for name, model in self.models.items()}
+        master_params_names = sum(
+            [[(name, n) for n, p in model.named_parameters() if p.requires_grad] for name, model in self.models.items()]
+        , [])
+        for i, (model_name, param_name) in enumerate(master_params_names):
+            state_dicts[model_name][param_name] = master_params[i]
+        return state_dicts
+
+    def _state_dicts_to_master_params(self, master_params, state_dicts):
+        """
+        Convert a state_dict to master params.
+        """
+        master_params_names = sum(
+            [[(name, n) for n, p in model.named_parameters() if p.requires_grad] for name, model in self.models.items()]
+        , [])
+        params = [state_dicts[name][param_name] for name, param_name in master_params_names]
+        if self.mix_precision_mode == 'inflat_all':
+            model_params_to_master_params(params, master_params)
+        else:
+            for i, param in enumerate(params):
+                master_params[i].data.copy_(param.data)
+
+    def load(self, load_dir, step=0):
+        """
+        Load a checkpoint.
+        Should be called by all processes.
+        """
+        if self.is_master:
+            print(f'\nLoading checkpoint from step {step}...', end='')
+            
+        model_ckpts = {}
+        for name, model in self.models.items():
+            model_ckpt = torch.load(read_file_dist(os.path.join(load_dir, 'ckpts', f'{name}_step{step:07d}.pt')), map_location=self.device, weights_only=True)
+            model_ckpts[name] = model_ckpt
+            model.load_state_dict(model_ckpt)
+        self._state_dicts_to_master_params(self.master_params, model_ckpts)
+        del model_ckpts
+
+        if self.is_master:
+            for i, ema_rate in enumerate(self.ema_rate):
+                ema_ckpts = {}
+                for name, model in self.models.items():
+                    ema_ckpt = torch.load(os.path.join(load_dir, 'ckpts', f'{name}_ema{ema_rate}_step{step:07d}.pt'), map_location=self.device, weights_only=True)
+                    ema_ckpts[name] = ema_ckpt
+                self._state_dicts_to_master_params(self.ema_params[i], ema_ckpts)
+                del ema_ckpts
+        
+        misc_ckpt = torch.load(read_file_dist(os.path.join(load_dir, 'ckpts', f'misc_step{step:07d}.pt')), map_location=torch.device('cpu'), weights_only=False)
+        self.optimizer.load_state_dict(misc_ckpt['optimizer'])
+        self.step = misc_ckpt['step']
+        self.data_sampler.load_state_dict(misc_ckpt['data_sampler'])
+        if self.mix_precision_mode == 'amp' and self.mix_precision_dtype == torch.float16:
+            self.scaler.load_state_dict(misc_ckpt['scaler'])
+        elif self.mix_precision_mode == 'inflat_all' and self.mix_precision_dtype == torch.float16:
+            self.log_scale = misc_ckpt['log_scale']
+        if self.lr_scheduler_config is not None:
+            self.lr_scheduler.load_state_dict(misc_ckpt['lr_scheduler'])
+        if self.elastic_controller_config is not None:
+            self.elastic_controller.load_state_dict(misc_ckpt['elastic_controller'])
+        if self.grad_clip is not None and not isinstance(self.grad_clip, float):
+            self.grad_clip.load_state_dict(misc_ckpt['grad_clip'])
+        del misc_ckpt
+
+        if self.world_size > 1:
+            dist.barrier()
+        if self.is_master:
+            print(' Done.')
+
+        if self.world_size > 1:
+            self.check_ddp()
+
+    def save(self, non_blocking=True):
+        """
+        Save a checkpoint.
+        Should be called only by the rank 0 process.
+        """
+        assert self.is_master, 'save() should be called only by the rank 0 process.'
+        print(f'\nSaving checkpoint at step {self.step}...', end='')
+        
+        model_ckpts = self._master_params_to_state_dicts(self.master_params)
+        for name, model_ckpt in model_ckpts.items():
+            model_ckpt = {k: v.cpu() for k, v in model_ckpt.items()}  # Move to CPU for saving
+            if non_blocking:
+                threading.Thread(
+                    target=torch.save,
+                    args=(model_ckpt, os.path.join(self.output_dir, 'ckpts', f'{name}_step{self.step:07d}.pt')),
+                ).start()
+            else:
+                torch.save(model_ckpt, os.path.join(self.output_dir, 'ckpts', f'{name}_step{self.step:07d}.pt'))
+        
+        for i, ema_rate in enumerate(self.ema_rate):
+            ema_ckpts = self._master_params_to_state_dicts(self.ema_params[i])
+            for name, ema_ckpt in ema_ckpts.items():
+                ema_ckpt = {k: v.cpu() for k, v in ema_ckpt.items()}  # Move to CPU for saving
+                if non_blocking:
+                    threading.Thread(
+                        target=torch.save,
+                        args=(ema_ckpt, os.path.join(self.output_dir, 'ckpts', f'{name}_ema{ema_rate}_step{self.step:07d}.pt')),
+                    ).start()
+                else:
+                    torch.save(ema_ckpt, os.path.join(self.output_dir, 'ckpts', f'{name}_ema{ema_rate}_step{self.step:07d}.pt'))
+
+        misc_ckpt = {
+            'optimizer': self.optimizer.state_dict(),
+            'step': self.step,
+            'data_sampler': self.data_sampler.state_dict(),
+        }
+        if self.mix_precision_mode == 'amp' and self.mix_precision_dtype == torch.float16:
+            misc_ckpt['scaler'] = self.scaler.state_dict()
+        elif self.mix_precision_mode == 'inflat_all' and self.mix_precision_dtype == torch.float16:
+            misc_ckpt['log_scale'] = self.log_scale
+        if self.lr_scheduler_config is not None:
+            misc_ckpt['lr_scheduler'] = self.lr_scheduler.state_dict()
+        if self.elastic_controller_config is not None:
+            misc_ckpt['elastic_controller'] = self.elastic_controller.state_dict()
+        if self.grad_clip is not None and not isinstance(self.grad_clip, float):
+            misc_ckpt['grad_clip'] = self.grad_clip.state_dict()
+        if non_blocking:
+            threading.Thread(
+                target=torch.save,
+                args=(misc_ckpt, os.path.join(self.output_dir, 'ckpts', f'misc_step{self.step:07d}.pt')),
+            ).start()
+        else:
+            torch.save(misc_ckpt, os.path.join(self.output_dir, 'ckpts', f'misc_step{self.step:07d}.pt'))
+        print(' Done.')
+
+    def _remap_checkpoint_keys(self, model_ckpt, model_state_dict):
+        """
+        Remap checkpoint keys to match model state dict.
+        
+        Handles structural changes like:
+        - cross_attn.xxx -> cross_attn.cross_attn_block.xxx (for ProjectAttention wrapper)
+        
+        Args:
+            model_ckpt: Checkpoint state dict
+            model_state_dict: Model state dict
+            
+        Returns:
+            Remapped checkpoint dict
+        """
+        remapped_ckpt = {}
+        remapped_count = 0
+        
+        for ckpt_key, ckpt_value in model_ckpt.items():
+            # Check if key exists directly
+            if ckpt_key in model_state_dict:
+                remapped_ckpt[ckpt_key] = ckpt_value
+                continue
+            
+            # Try remapping: cross_attn.xxx -> cross_attn.cross_attn_block.xxx
+            # This handles the case when cross_attn is wrapped by ProjectAttention
+            if '.cross_attn.' in ckpt_key:
+                # Split at .cross_attn.
+                parts = ckpt_key.split('.cross_attn.')
+                if len(parts) == 2:
+                    new_key = f'{parts[0]}.cross_attn.cross_attn_block.{parts[1]}'
+                    if new_key in model_state_dict:
+                        remapped_ckpt[new_key] = ckpt_value
+                        remapped_count += 1
+                        continue
+            
+            # Key not remapped, keep original (will be handled by missing key logic)
+            remapped_ckpt[ckpt_key] = ckpt_value
+        
+        if remapped_count > 0 and self.is_master:
+            print(f'Info: Remapped {remapped_count} cross_attn keys to cross_attn.cross_attn_block')
+        
+        return remapped_ckpt
+
+    def finetune_from(self, finetune_ckpt):
+        """
+        Finetune from a checkpoint.
+        Should be called by all processes.
+        """
+        # Allow missing keys (e.g., register_buffer parameters)
+        ALLOWED_MISSING_KEYS = {'rope_phases'}
+        
+        if self.is_master:
+            print('\nFinetuning from:')
+            for name, path in finetune_ckpt.items():
+                print(f'  - {name}: {path}')
+        
+        model_ckpts = {}
+        for name, model in self.models.items():
+            model_state_dict = model.state_dict()
+            if name in finetune_ckpt:
+                model_ckpt = torch.load(read_file_dist(finetune_ckpt[name]), map_location=self.device, weights_only=True)
+                
+                # Remap checkpoint keys to handle structural changes (e.g., ProjectAttention wrapper)
+                model_ckpt = self._remap_checkpoint_keys(model_ckpt, model_state_dict)
+                
+                # Check extra keys (in ckpt but not in model)
+                for k, v in model_ckpt.items():
+                    if k not in model_state_dict:
+                        if self.is_master:
+                            print(f'Warning: {k} not found in model_state_dict, skipped.')
+                        model_ckpt[k] = None
+                    elif model_ckpt[k].shape != model_state_dict[k].shape:
+                        if self.is_master:
+                            print(f'Warning: {k} shape mismatch, {model_ckpt[k].shape} vs {model_state_dict[k].shape}, skipped.')
+                        model_ckpt[k] = model_state_dict[k]
+                model_ckpt = {k: v for k, v in model_ckpt.items() if v is not None}
+                
+                # Check missing keys (in model but not in ckpt)
+                missing_keys = set(model_state_dict.keys()) - set(model_ckpt.keys())
+                unexpected_missing = missing_keys - ALLOWED_MISSING_KEYS
+                if unexpected_missing and self.is_master:
+                    print(f'Error: Missing keys in checkpoint: {unexpected_missing}')
+                    raise RuntimeError(f'Missing keys in checkpoint: {unexpected_missing}')
+                if missing_keys & ALLOWED_MISSING_KEYS and self.is_master:
+                    print(f'Info: Using model initialized values for: {missing_keys & ALLOWED_MISSING_KEYS}')
+                
+                # Fill in missing keys (using model initialized values)
+                for k in missing_keys:
+                    model_ckpt[k] = model_state_dict[k]
+                
+                model_ckpts[name] = model_ckpt
+                model.load_state_dict(model_ckpt)
+            else:
+                if self.is_master:
+                    print(f'Warning: {name} not found in finetune_ckpt, skipped.')
+                model_ckpts[name] = model_state_dict
+        self._state_dicts_to_master_params(self.master_params, model_ckpts)
+        if self.is_master:
+            for i, ema_rate in enumerate(self.ema_rate):
+                self._state_dicts_to_master_params(self.ema_params[i], model_ckpts)
+        del model_ckpts
+
+        if self.world_size > 1:
+            dist.barrier()
+        if self.is_master:
+            print('Done.')
+
+        if self.world_size > 1:
+            self.check_ddp()
+
+    @abstractmethod
+    def run_snapshot(self, num_samples, batch_size=4, verbose=False, **kwargs):
+        """
+        Run a snapshot of the model.
+        """
+        pass
+
+    @torch.no_grad()
+    def visualize_sample(self, sample):
+        """
+        Convert a sample to an image.
+        """
+        if hasattr(self.dataset, 'visualize_sample'):
+            return self.dataset.visualize_sample(sample)
+        else:
+            return sample
+
+    @torch.no_grad()
+    def snapshot_dataset(self, num_samples=100, batch_size=4):
+        """
+        Sample images from the dataset.
+        """
+        dataloader = torch.utils.data.DataLoader(
+            self.dataset,
+            batch_size=batch_size,
+            num_workers=0,
+            shuffle=True,
+            collate_fn=self.dataset.collate_fn if hasattr(self.dataset, 'collate_fn') else None,
+        )
+        save_cfg = {}
+        for i in range(0, num_samples, batch_size):
+            data = next(iter(dataloader))
+            data = {k: v[:min(num_samples - i, batch_size)] for k, v in data.items()}
+            data = recursive_to_device(data, self.device)
+            try:
+                vis = self.visualize_sample(data)
+            except (RuntimeError, Exception) as e:
+                print(f'\033[93m[WARN] snapshot_dataset visualize_sample failed (batch {i}), skipping: {e}\033[0m')
+                torch.cuda.empty_cache()
+                continue
+            if isinstance(vis, dict):
+                for k, v in vis.items():
+                    if f'dataset_{k}' not in save_cfg:
+                        save_cfg[f'dataset_{k}'] = []
+                    save_cfg[f'dataset_{k}'].append(v)
+            else:
+                if 'dataset' not in save_cfg:
+                    save_cfg['dataset'] = []
+                save_cfg['dataset'].append(vis)
+        for name, image in save_cfg.items():
+            utils.save_image(
+                torch.cat(image, dim=0),
+                os.path.join(self.output_dir, 'samples', f'{name}.jpg'),
+                nrow=int(np.sqrt(num_samples)),
+                normalize=True,
+                value_range=self.dataset.value_range,
+            )
+
+    @torch.no_grad()
+    def snapshot(self, suffix=None, num_samples=64, batch_size=4, verbose=False):
+        """
+        Sample images from the model.
+        NOTE: When num_samples >= 4, this function should be called by all processes.
+              When num_samples < 4, only master runs snapshot (other ranks skip via barrier).
+        """
+        # Free cached GPU memory before snapshot to avoid OOM / illegal address errors
+        import gc
+        gc.collect()
+        torch.cuda.empty_cache()
+
+        if self.is_master:
+            print(f'\nSampling {num_samples} images...', end='')
+
+        if suffix is None:
+            suffix = f'step{self.step:07d}'
+
+        # When num_samples < 4, only master runs snapshot to avoid multi-rank gather issues
+        master_only = num_samples < 4
+        
+        sample_metadata = None  # Will hold list of "dataset_name/sha256" strings
+
+        if master_only and self.world_size > 1:
+            if not self.is_master:
+                # Non-master ranks just wait at barrier
+                dist.barrier()
+                return
+
+            # Master runs snapshot alone
+            amp_context = partial(torch.autocast, device_type='cuda', dtype=self.mix_precision_dtype) if self.mix_precision_mode == 'amp' else nullcontext
+            with amp_context():
+                samples = self.run_snapshot(num_samples, batch_size=batch_size, verbose=verbose)
+
+            # Extract metadata before preprocessing
+            sample_metadata = samples.pop('_metadata', None)
+
+            # Free GPU memory after sampling, before decode + render
+            torch.cuda.empty_cache()
+
+            # Preprocess images
+            for key in list(samples.keys()):
+                if samples[key]['type'] == 'sample':
+                    try:
+                        vis = self.visualize_sample(samples[key]['value'])
+                    except RuntimeError as e:
+                        print(f"[Snapshot] WARNING: visualize_sample failed for '{key}': {e}")
+                        # Reset CUDA error state and skip this sample
+                        try:
+                            torch.cuda.synchronize()
+                        except RuntimeError:
+                            pass
+                        torch.cuda.empty_cache()
+                        del samples[key]
+                        continue
+                    if isinstance(vis, dict):
+                        for k, v in vis.items():
+                            samples[f'{key}_{k}'] = {'value': v, 'type': 'image'}
+                        del samples[key]
+                    else:
+                        samples[key] = {'value': vis, 'type': 'image'}
+
+            # No gather needed, master already has all samples
+            dist.barrier()
+        else:
+            # Distribute sampling across all ranks
+            num_samples_per_process = int(np.ceil(num_samples / self.world_size))
+            amp_context = partial(torch.autocast, device_type='cuda', dtype=self.mix_precision_dtype) if self.mix_precision_mode == 'amp' else nullcontext
+            
+            with amp_context():
+                samples = self.run_snapshot(num_samples_per_process, batch_size=batch_size, verbose=verbose)
+
+            # Extract metadata before preprocessing
+            local_metadata = samples.pop('_metadata', None)
+
+            # Free GPU memory after sampling, before decode + render
+            torch.cuda.empty_cache()
+
+            # Preprocess images
+            for key in list(samples.keys()):
+                if samples[key]['type'] == 'sample':
+                    try:
+                        vis = self.visualize_sample(samples[key]['value'])
+                    except RuntimeError as e:
+                        print(f"[Snapshot] WARNING: visualize_sample failed for '{key}': {e}")
+                        torch.cuda.synchronize()
+                        del samples[key]
+                        continue
+                    if isinstance(vis, dict):
+                        for k, v in vis.items():
+                            samples[f'{key}_{k}'] = {'value': v, 'type': 'image'}
+                        del samples[key]
+                    else:
+                        samples[key] = {'value': vis, 'type': 'image'}
+
+            # Gather results
+            if self.world_size > 1:
+                for key in samples.keys():
+                    samples[key]['value'] = samples[key]['value'].contiguous()
+                    if self.is_master:
+                        all_images = [torch.empty_like(samples[key]['value']) for _ in range(self.world_size)]
+                    else:
+                        all_images = []
+                    dist.gather(samples[key]['value'], all_images, dst=0)
+                    if self.is_master:
+                        samples[key]['value'] = torch.cat(all_images, dim=0)[:num_samples]
+                
+                # Gather metadata across ranks
+                if local_metadata is not None:
+                    all_metadata = [None] * self.world_size
+                    dist.all_gather_object(all_metadata, local_metadata)
+                    if self.is_master:
+                        sample_metadata = sum(all_metadata, [])[:num_samples]
+                else:
+                    sample_metadata = None
+            else:
+                sample_metadata = local_metadata
+
+        # Save images
+        if self.is_master:
+            os.makedirs(os.path.join(self.output_dir, 'samples', suffix), exist_ok=True)
+            wandb_images = {}  # Collect images for wandb logging
+            nrow = int(np.sqrt(num_samples))
+            vr = self.dataset.value_range
+            
+            # Build metadata caption string for wandb
+            metadata_caption = ''
+            if sample_metadata:
+                metadata_caption = '\n' + ' | '.join(sample_metadata)
+                # Also save metadata to file
+                with open(os.path.join(self.output_dir, 'samples', suffix, 'metadata.txt'), 'w') as f:
+                    for i, m in enumerate(sample_metadata):
+                        f.write(f'{i}: {m}\n')
+
+            # Helper: make a normalized grid tensor from a batch of images
+            def _make_grid(tensor):
+                return utils.make_grid(tensor, nrow=nrow, normalize=True, value_range=vr)
+
+            # Helper: resize grid to target height (keep aspect ratio)
+            def _resize_to_height(grid, target_h):
+                import torch.nn.functional as F
+                _, h, w = grid.shape
+                if h == target_h:
+                    return grid
+                target_w = int(round(w * target_h / h))
+                return F.interpolate(grid.unsqueeze(0), size=(target_h, target_w), mode='bilinear', align_corners=False).squeeze(0)
+
+            # --- Save individual images (original behavior) ---
+            for key in samples.keys():
+                if samples[key]['type'] == 'image':
+                    image_path = os.path.join(self.output_dir, 'samples', suffix, f'{key}_{suffix}.jpg')
+                    utils.save_image(
+                        samples[key]['value'],
+                        image_path,
+                        nrow=nrow,
+                        normalize=True,
+                        value_range=vr,
+                    )
+                    # Collect for wandb
+                    if self.wandb_run is not None:
+                        grid = _make_grid(samples[key]['value'])
+                        grid_np = grid.permute(1, 2, 0).cpu().numpy()
+                        grid_np = (grid_np * 255).clip(0, 255).astype(np.uint8)
+                        wandb_images[f'samples/{key}'] = wandb.Image(grid_np, caption=f'{key} at step {self.step}{metadata_caption}')
+                elif samples[key]['type'] == 'number':
+                    val_min = samples[key]['value'].min()
+                    val_max = samples[key]['value'].max()
+                    images = (samples[key]['value'] - val_min) / (val_max - val_min)
+                    images = utils.make_grid(
+                        images,
+                        nrow=nrow,
+                        normalize=False,
+                    )
+                    save_image_with_notes(
+                        images,
+                        os.path.join(self.output_dir, 'samples', suffix, f'{key}_{suffix}.jpg'),
+                        notes=f'{key} min: {val_min}, max: {val_max}',
+                    )
+
+            # --- Save combined images ---
+            sample_keys = set(samples.keys())
+
+            # Combined 1: image + sample_gt_view + sample_gt_gt_view (shape)
+            #             image + sample_gt_view_{attr} + sample_gt_gt_view_{attr} (tex, per attribute)
+            # Detect gt_view attribute suffixes from sample keys
+            gt_view_attrs = set()
+            for k in sample_keys:
+                if k.startswith('sample_gt_view_'):
+                    attr = k[len('sample_gt_view_'):]
+                    gt_view_attrs.add(attr)
+            
+            if gt_view_attrs:
+                # Tex mode: generate combined view for each PBR attribute
+                for attr in sorted(gt_view_attrs):
+                    combo1_keys = ['image', f'sample_gt_view_{attr}', f'sample_gt_gt_view_{attr}']
+                    combo1_present = [k for k in combo1_keys if k in sample_keys and samples[k]['type'] == 'image']
+                    if len(combo1_present) >= 2:
+                        grids = [_make_grid(samples[k]['value']) for k in combo1_present]
+                        target_h = max(g.shape[1] for g in grids)
+                        grids = [_resize_to_height(g, target_h) for g in grids]
+                        combined = torch.cat(grids, dim=2)
+                        combined_path = os.path.join(self.output_dir, 'samples', suffix, f'combined_views_{attr}_{suffix}.jpg')
+                        utils.save_image(combined, combined_path, normalize=False)
+                        if self.wandb_run is not None:
+                            grid_np = combined.permute(1, 2, 0).cpu().numpy()
+                            grid_np = (grid_np * 255).clip(0, 255).astype(np.uint8)
+                            label = ' | '.join(combo1_present)
+                            wandb_images[f'samples/combined_views_{attr}'] = wandb.Image(grid_np, caption=f'{label} at step {self.step}{metadata_caption}')
+            else:
+                # Shape mode: single gt_view
+                combo1_keys = ['image', 'sample_gt_view', 'sample_gt_gt_view']
+                combo1_present = [k for k in combo1_keys if k in sample_keys and samples[k]['type'] == 'image']
+                if len(combo1_present) >= 2:
+                    grids = [_make_grid(samples[k]['value']) for k in combo1_present]
+                    target_h = max(g.shape[1] for g in grids)
+                    grids = [_resize_to_height(g, target_h) for g in grids]
+                    combined = torch.cat(grids, dim=2)
+                    combined_path = os.path.join(self.output_dir, 'samples', suffix, f'combined_views_{suffix}.jpg')
+                    utils.save_image(combined, combined_path, normalize=False)
+                    if self.wandb_run is not None:
+                        grid_np = combined.permute(1, 2, 0).cpu().numpy()
+                        grid_np = (grid_np * 255).clip(0, 255).astype(np.uint8)
+                        label = ' | '.join(combo1_present)
+                        wandb_images[f'samples/combined_views'] = wandb.Image(grid_np, caption=f'{label} at step {self.step}{metadata_caption}')
+
+            # Combined 2: sample_multiview + sample_gt_multiview
+            combo2_keys = ['sample_multiview', 'sample_gt_multiview']
+            combo2_present = [k for k in combo2_keys if k in sample_keys and samples[k]['type'] == 'image']
+            if len(combo2_present) >= 2:
+                grids = [_make_grid(samples[k]['value']) for k in combo2_present]
+                target_h = max(g.shape[1] for g in grids)
+                grids = [_resize_to_height(g, target_h) for g in grids]
+                combined = torch.cat(grids, dim=2)  # concatenate along width
+                combined_path = os.path.join(self.output_dir, 'samples', suffix, f'combined_multiview_{suffix}.jpg')
+                utils.save_image(combined, combined_path, normalize=False)
+                if self.wandb_run is not None:
+                    grid_np = combined.permute(1, 2, 0).cpu().numpy()
+                    grid_np = (grid_np * 255).clip(0, 255).astype(np.uint8)
+                    label = ' | '.join(combo2_present)
+                    wandb_images[f'samples/combined_multiview'] = wandb.Image(grid_np, caption=f'{label} at step {self.step}{metadata_caption}')
+
+            # Log images to wandb
+            if self.wandb_run is not None and wandb_images:
+                self.wandb_run.log(wandb_images, step=self.step)
+
+        if self.is_master:
+            print(' Done.')
+
+    def update_ema(self):
+        """
+        Update exponential moving average.
+        Should only be called by the rank 0 process.
+        """
+        assert self.is_master, 'update_ema() should be called only by the rank 0 process.'
+        for i, ema_rate in enumerate(self.ema_rate):
+            for master_param, ema_param in zip(self.master_params, self.ema_params[i]):
+                ema_param.detach().mul_(ema_rate).add_(master_param, alpha=1.0 - ema_rate)
+
+    def check_ddp(self):
+        """
+        Check if DDP is working properly.
+        Should be called by all process.
+        """
+        if self.is_master:
+            print('\nPerforming DDP check...')
+
+        if self.is_master:
+            print('Checking if parameters are consistent across processes...')
+        dist.barrier()
+        try:
+            for p in self.master_params:
+                # split to avoid OOM
+                for i in range(0, p.numel(), 10000000):
+                    sub_size = min(10000000, p.numel() - i)
+                    sub_p = p.detach().view(-1)[i:i+sub_size]
+                    # gather from all processes
+                    sub_p_gather = [torch.empty_like(sub_p) for _ in range(self.world_size)]
+                    dist.all_gather(sub_p_gather, sub_p)
+                    # check if equal
+                    assert all([torch.equal(sub_p, sub_p_gather[i]) for i in range(self.world_size)]), 'parameters are not consistent across processes'
+        except AssertionError as e:
+            if self.is_master:
+                print(f'\n\033[91mError: {e}\033[0m')
+                print('DDP check failed.')
+            raise e
+
+        dist.barrier()
+        if self.is_master:
+            print('Done.')
+
+    def _verify_gradient_sync(self):
+        """
+        Verify that DDP gradient synchronization is working correctly.
+        DDP's backward automatically performs all_reduce on gradients; after sync all ranks should have identical gradients.
+        
+        Verification method:
+        1. Compute total gradient norm across all parameters
+        2. Gather gradient norms from all ranks
+        3. If DDP sync is working, all ranks should have identical gradient norms
+        4. If not synced, gradient norms will differ (since each rank processes different data)
+        """
+        # Compute total gradient norm on this rank
+        total_grad_norm_sq = 0.0
+        grad_count = 0
+        for p in self.model_params:
+            if p.grad is not None:
+                total_grad_norm_sq += p.grad.detach().float().norm().item() ** 2
+                grad_count += 1
+        
+        if grad_count == 0:
+            return
+        
+        local_grad_norm = total_grad_norm_sq ** 0.5
+        
+        # Ensure all processes reach the same point
+        dist.barrier()
+        
+        # Gather gradient norms from all ranks
+        grad_norm_tensor = torch.tensor([local_grad_norm], dtype=torch.float64, device=self.device)
+        all_grad_norms = [torch.zeros(1, dtype=torch.float64, device=self.device) for _ in range(self.world_size)]
+        dist.all_gather(all_grad_norms, grad_norm_tensor)
+        all_grad_norms = [g.item() for g in all_grad_norms]
+        
+        # Verify all ranks have the same gradient norm (relative error tolerance: 0.1%)
+        ref_norm = all_grad_norms[0]
+        if ref_norm > 0:
+            is_synced = all(abs(g - ref_norm) / ref_norm < 1e-3 for g in all_grad_norms)
+        else:
+            is_synced = all(abs(g) < 1e-10 for g in all_grad_norms)
+        
+        if self.is_master:
+            print(f'\n{"="*60}')
+            print(f'[Step {self.step}] DDP Gradient Sync Verification:')
+            for i, g in enumerate(all_grad_norms):
+                print(f'  Rank {i} grad_norm: {g:.8f}')
+            if is_synced:
+                print(f'  \033[92m✓ PASS: All gradients are synchronized!\033[0m')
+            else:
+                max_diff = max(abs(g - ref_norm) for g in all_grad_norms)
+                print(f'  \033[91m✗ FAIL: Gradients are NOT synchronized! Max diff: {max_diff:.8f}\033[0m')
+            print(f'{"="*60}\n')
+
+    @abstractmethod
+    def training_losses(**mb_data):
+        """
+        Compute training losses.
+        """
+        pass
+
+    def load_data(self):
+        """
+        Load data.
+        """
+        if self.prefetch_data:
+            if self._data_prefetched is None:
+                self._data_prefetched = recursive_to_device(next(self.data_iterator), self.device, non_blocking=True)
+            data = self._data_prefetched
+            self._data_prefetched = recursive_to_device(next(self.data_iterator), self.device, non_blocking=True)
+        else:
+            data = recursive_to_device(next(self.data_iterator), self.device, non_blocking=True)
+        
+        # if the data is a dict, we need to split it into multiple dicts with batch_size_per_gpu
+        if isinstance(data, dict):
+            if self.batch_split == 1:
+                data_list = [data]
+            else:
+                batch_size = list(data.values())[0].shape[0]
+                data_list = [
+                    {k: v[i * batch_size // self.batch_split:(i + 1) * batch_size // self.batch_split] for k, v in data.items()}
+                    for i in range(self.batch_split)
+                ]
+        elif isinstance(data, list):
+            data_list = data
+        else:
+            raise ValueError('Data must be a dict or a list of dicts.')
+        
+        return data_list
+
+    def run_step(self, data_list):
+        """
+        Run a training step.
+        """
+        step_log = {'loss': {}, 'status': {}}
+        amp_context = partial(torch.autocast, device_type='cuda', dtype=self.mix_precision_dtype) if self.mix_precision_mode == 'amp' else nullcontext
+        elastic_controller_context = self.elastic_controller.record if self.elastic_controller_config is not None else nullcontext
+
+        # Train
+        losses = []
+        statuses = []
+        elastic_controller_logs = []
+        zero_grad(self.model_params)
+        for i, mb_data in enumerate(data_list):
+            ## sync at the end of each batch split
+            sync_contexts = [self.training_models[name].no_sync for name in self.training_models] if i != len(data_list) - 1 and self.world_size > 1 else [nullcontext]
+            with nested_contexts(*sync_contexts), elastic_controller_context():
+                with amp_context():
+                    loss, status = self.training_losses(**mb_data)
+                    l = loss['loss'] / len(data_list)
+                    
+                    # DEBUG: Print loss for each rank
+                    if self.debug:
+                        print(f'[Rank {self.rank}/{self.world_size}] Step {self.step} batch {i}: loss={loss["loss"].item():.6f}')
+                    
+                ## backward
+                if self.mix_precision_mode == 'amp' and self.mix_precision_dtype == torch.float16:
+                    self.scaler.scale(l).backward()
+                elif self.mix_precision_mode == 'inflat_all' and self.mix_precision_dtype == torch.float16:
+                    scaled_l = l * (2 ** self.log_scale)
+                    scaled_l.backward()
+                else:
+                    l.backward()
+            ## log
+            losses.append(dict_foreach(loss, lambda x: x.item() if isinstance(x, torch.Tensor) else x))
+            statuses.append(dict_foreach(status, lambda x: x.item() if isinstance(x, torch.Tensor) else x))
+            if self.elastic_controller_config is not None:
+                elastic_controller_logs.append(self.elastic_controller.log())
+        
+        # ============================================================
+        # DEBUG: Verify DDP gradient synchronization
+        # Check if gradients are consistent across ranks after backward
+        # DDP automatically all_reduces gradients during the last batch_split's backward
+        # After sync, all ranks should have identical gradients
+        # ============================================================
+        if self.debug and self.world_size > 1:
+            self._verify_gradient_sync()
+        
+        ## gradient clip
+        if self.grad_clip is not None:
+            if self.mix_precision_mode == 'amp' and self.mix_precision_dtype == torch.float16:
+                self.scaler.unscale_(self.optimizer)
+            elif self.mix_precision_mode == 'inflat_all':
+                model_grads_to_master_grads(self.model_params, self.master_params)
+                if self.mix_precision_dtype == torch.float16:
+                    self.master_params[0].grad.mul_(1.0 / (2 ** self.log_scale))
+            if isinstance(self.grad_clip, float):
+                grad_norm = torch.nn.utils.clip_grad_norm_(self.master_params, self.grad_clip)
+            else:
+                grad_norm = self.grad_clip(self.master_params)
+            if torch.isfinite(grad_norm):
+                statuses[-1]['grad_norm'] = grad_norm.item()
+        ## step
+        if self.mix_precision_mode == 'amp' and self.mix_precision_dtype == torch.float16:
+            prev_scale = self.scaler.get_scale()
+            self.scaler.step(self.optimizer)
+            self.scaler.update()
+        elif self.mix_precision_mode == 'inflat_all':
+            if self.mix_precision_dtype == torch.float16:
+                prev_scale = 2 ** self.log_scale
+                if not any(not p.grad.isfinite().all() for p in self.model_params):
+                    if self.grad_clip is None:
+                        model_grads_to_master_grads(self.model_params, self.master_params)
+                        self.master_params[0].grad.mul_(1.0 / (2 ** self.log_scale))
+                    self.optimizer.step()
+                    master_params_to_model_params(self.model_params, self.master_params)
+                    self.log_scale += self.fp16_scale_growth
+                else:
+                    self.log_scale -= 1
+            else:
+                prev_scale = 1.0
+                if self.grad_clip is None:
+                    model_grads_to_master_grads(self.model_params, self.master_params)
+                if not any(not p.grad.isfinite().all() for p in self.master_params):
+                    self.optimizer.step()
+                    master_params_to_model_params(self.model_params, self.master_params)
+                else:
+                    print('\n\033[93mWarning: NaN detected in gradients. Skipping update.\033[0m')
+        else:
+            prev_scale = 1.0
+            if not any(not p.grad.isfinite().all() for p in self.model_params):
+                self.optimizer.step()
+            else:
+                print('\n\033[93mWarning: NaN detected in gradients. Skipping update.\033[0m') 
+        ## adjust learning rate
+        if self.lr_scheduler_config is not None:
+            statuses[-1]['lr'] = self.lr_scheduler.get_last_lr()[0]
+            self.lr_scheduler.step()
+
+        # Logs
+        step_log['loss'] = dict_reduce(losses, lambda x: np.mean(x))
+        step_log['status'] = dict_reduce(statuses, lambda x: np.mean(x), special_func={'min': lambda x: np.min(x), 'max': lambda x: np.max(x)})
+        if self.elastic_controller_config is not None:
+            step_log['elastic'] = dict_reduce(elastic_controller_logs, lambda x: np.mean(x))
+        if self.grad_clip is not None:
+            step_log['grad_clip'] = self.grad_clip if isinstance(self.grad_clip, float) else self.grad_clip.log()
+            
+        # Check grad and norm of each param
+        if self.log_param_stats:
+            param_norms = {}
+            param_grads = {}
+            for model_name, model in self.models.items():
+                for name, param in model.named_parameters():
+                    if param.requires_grad:
+                        param_norms[f'{model_name}.{name}'] = param.norm().item()
+                        if param.grad is not None and torch.isfinite(param.grad).all():
+                            param_grads[f'{model_name}.{name}'] = param.grad.norm().item() / prev_scale
+            step_log['param_norms'] = param_norms
+            step_log['param_grads'] = param_grads
+
+        # Update exponential moving average
+        if self.is_master:
+            self.update_ema()
+
+        return step_log
+
+    def save_logs(self):
+        log_str = '\n'.join([
+            f'{step}: {json.dumps(dict_foreach(log, lambda x: float(x)))}' for step, log in self.log
+        ])
+        
+        # Accumulate logs in memory and overwrite file each time (S3 FUSE does not support append)
+        if not hasattr(self, '_log_buffer'):
+            self._log_buffer = []
+        self._log_buffer.append(log_str)
+        try:
+            with open(self._log_file, 'w') as log_file:
+                log_file.write('\n'.join(self._log_buffer) + '\n')
+        except Exception as e:
+            print(f'\033[93m[WARN] Failed to write log file: {e}\033[0m')
+
+        # show with mlflow
+        log_show = [l for _, l in self.log if not dict_any(l, lambda x: np.isnan(x))]
+        log_show = dict_reduce(log_show, lambda x: np.mean(x))
+        log_show = dict_flatten(log_show, sep='/')
+        if self.writer is not None:
+            for key, value in log_show.items():
+                self.writer.add_scalar(key, value, self.step)
+        
+        # Log to wandb
+        if self.wandb_run is not None:
+            wandb_log = {key: value for key, value in log_show.items()}
+            wandb_log['step'] = self.step
+            self.wandb_run.log(wandb_log, step=self.step)
+        
+        self.log = []
+        
+    def check_abort(self):
+        """
+        Check if training should be aborted due to certain conditions.
+        """
+        # 1. If log_scale in inflat_all mode is less than 0
+        if self.mix_precision_dtype == torch.float16 and \
+           self.mix_precision_mode == 'inflat_all' and \
+           self.log_scale < 0:
+            if self.is_master:
+                print ('\n\n\033[91m')
+                print (f'ABORT: log_scale in inflat_all mode is less than 0 at step {self.step}.')
+                print ('This indicates that the model is diverging. You should look into the model and the data.')
+                print ('\033[0m')
+                self.save(non_blocking=False)
+                self.save_logs()
+            if self.world_size > 1:
+                dist.barrier()
+            raise ValueError('ABORT: log_scale in inflat_all mode is less than 0.')
+
+    def run(self):
+        """
+        Run training.
+        """
+        if self.is_master:
+            print('\nStarting training...')
+            if self.i_sample != -1:
+                try:
+                    self.snapshot_dataset(num_samples=self.snapshot_num_samples, batch_size=self.snapshot_batch_size)
+                except (RuntimeError, Exception) as e:
+                    print(f'\033[93m[WARN] snapshot_dataset failed, skipping: {e}\033[0m')
+                    torch.cuda.empty_cache()
+            else:
+                print('[INFO] i_sample=-1, all snapshots disabled.')
+        if self.i_sample != -1:
+            if self.step == 0:
+                try:
+                    self.snapshot(suffix='init', num_samples=self.snapshot_num_samples, batch_size=self.snapshot_batch_size)
+                except (RuntimeError, Exception) as e:
+                    print(f'\033[93m[WARN] snapshot (init) failed, skipping: {e}\033[0m')
+                    torch.cuda.empty_cache()
+            else: # resume
+                try:
+                    self.snapshot(suffix=f'resume_step{self.step:07d}', num_samples=self.snapshot_num_samples, batch_size=self.snapshot_batch_size)
+                except (RuntimeError, Exception) as e:
+                    print(f'\033[93m[WARN] snapshot (resume) failed, skipping: {e}\033[0m')
+                    torch.cuda.empty_cache()
+
+        time_last_print = 0.0
+        time_elapsed = 0.0
+        while self.step < self.max_steps:
+            time_start = time.time()
+
+            data_list = self.load_data()
+            step_log = self.run_step(data_list)
+
+            time_end = time.time()
+            time_elapsed += time_end - time_start
+
+            self.step += 1
+
+            # Print progress
+            if self.is_master and self.step % self.i_print == 0:
+                speed = self.i_print / (time_elapsed - time_last_print) * 3600
+                columns = [
+                    f'Step: {self.step}/{self.max_steps} ({self.step / self.max_steps * 100:.2f}%)',
+                    f'Elapsed: {time_elapsed / 3600:.2f} h',
+                    f'Speed: {speed:.2f} steps/h',
+                    f'ETA: {(self.max_steps - self.step) / speed:.2f} h',
+                ]
+                print(' | '.join([c.ljust(25) for c in columns]), flush=True)
+                time_last_print = time_elapsed
+
+            # Check ddp
+            if self.parallel_mode == 'ddp' and self.world_size > 1 and self.i_ddpcheck is not None and self.step % self.i_ddpcheck == 0:
+                self.check_ddp()
+
+            # Sample images
+            if self.i_sample != -1 and self.step % self.i_sample == 0:
+                try:
+                    self.snapshot(num_samples=self.snapshot_num_samples, batch_size=self.snapshot_batch_size)
+                except (RuntimeError, Exception) as e:
+                    if self.is_master:
+                        print(f'\033[93m[WARN] snapshot at step {self.step} failed, skipping: {e}\033[0m')
+                    try:
+                        torch.cuda.empty_cache()
+                    except Exception:
+                        pass
+
+            if self.is_master:
+                self.log.append((self.step, {}))
+
+                # Log time
+                self.log[-1][1]['time'] = {
+                    'step': time_end - time_start,
+                    'elapsed': time_elapsed,
+                }
+
+                # Log losses
+                if step_log is not None:
+                    self.log[-1][1].update(step_log)
+
+                # Log scale
+                if self.mix_precision_dtype == torch.float16:
+                    if self.mix_precision_mode == 'amp':
+                        self.log[-1][1]['scale'] = self.scaler.get_scale()
+                    elif self.mix_precision_mode == 'inflat_all':
+                        self.log[-1][1]['log_scale'] = self.log_scale
+
+                # Save log
+                if self.step % self.i_log == 0:
+                    self.save_logs()
+
+                # Save checkpoint
+                if self.step % self.i_save == 0:
+                    self.save()
+                    
+            # Check abort
+            self.check_abort()
+
+        if self.i_sample != -1:
+            try:
+                self.snapshot(suffix='final', num_samples=self.snapshot_num_samples, batch_size=self.snapshot_batch_size)
+            except (RuntimeError, Exception) as e:
+                if self.is_master:
+                    print(f'\033[93m[WARN] snapshot (final) failed, skipping: {e}\033[0m')
+            torch.cuda.empty_cache()
+        if self.world_size > 1:
+            dist.barrier()
+        if self.is_master:
+            self.writer.close()
+            print('Training finished.')
+            
+    def profile(self, wait=2, warmup=3, active=5):
+        """
+        Profile the training loop.
+        """
+        with torch.profiler.profile(
+            schedule=torch.profiler.schedule(wait=wait, warmup=warmup, active=active, repeat=1),
+            on_trace_ready=torch.profiler.tensorboard_trace_handler(os.path.join(self.output_dir, 'profile')),
+            profile_memory=True,
+            with_stack=True,
+        ) as prof:
+            for _ in range(wait + warmup + active):
+                self.run_step()
+                prof.step()
diff --git a/pixal3d/trainers/flow_matching/flow_matching.py b/pixal3d/trainers/flow_matching/flow_matching.py
new file mode 100644
index 0000000..1ab1bea
--- /dev/null
+++ b/pixal3d/trainers/flow_matching/flow_matching.py
@@ -0,0 +1,655 @@
+from typing import *
+import copy
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+import numpy as np
+from easydict import EasyDict as edict
+
+from ..basic import BasicTrainer
+from ...pipelines import samplers 
+from ...utils.general_utils import dict_reduce
+from .mixins.classifier_free_guidance import ClassifierFreeGuidanceMixin
+from .mixins.text_conditioned import TextConditionedMixin
+from .mixins.image_conditioned import ImageConditionedMixin
+from .mixins.image_conditioned_proj import ImageConditionedProjMixin
+
+
+class FlowMatchingTrainer(BasicTrainer):
+    """
+    Trainer for diffusion model with flow matching objective.
+    
+    Args:
+        models (dict[str, nn.Module]): Models to train.
+        dataset (torch.utils.data.Dataset): Dataset.
+        output_dir (str): Output directory.
+        load_dir (str): Load directory.
+        step (int): Step to load.
+        batch_size (int): Batch size.
+        batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
+        batch_split (int): Split batch with gradient accumulation.
+        max_steps (int): Max steps.
+        optimizer (dict): Optimizer config.
+        lr_scheduler (dict): Learning rate scheduler config.
+        elastic (dict): Elastic memory management config.
+        grad_clip (float or dict): Gradient clip config.
+        ema_rate (float or list): Exponential moving average rates.
+        fp16_mode (str): FP16 mode.
+            - None: No FP16.
+            - 'inflat_all': Hold a inflated fp32 master param for all params.
+            - 'amp': Automatic mixed precision.
+        fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
+        finetune_ckpt (dict): Finetune checkpoint.
+        log_param_stats (bool): Log parameter stats.
+        i_print (int): Print interval.
+        i_log (int): Log interval.
+        i_sample (int): Sample interval.
+        i_save (int): Save interval.
+        i_ddpcheck (int): DDP check interval.
+
+        t_schedule (dict): Time schedule for flow matching.
+        sigma_min (float): Minimum noise level.
+    """
+    def __init__(
+        self,
+        *args,
+        t_schedule: dict = {
+            'name': 'logitNormal',
+            'args': {
+                'mean': 0.0,
+                'std': 1.0,
+            }
+        },
+        sigma_min: float = 1e-5,
+        **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        self.t_schedule = t_schedule
+        self.sigma_min = sigma_min
+
+    def diffuse(self, x_0: torch.Tensor, t: torch.Tensor, noise: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Diffuse the data for a given number of diffusion steps.
+        In other words, sample from q(x_t | x_0).
+
+        Args:
+            x_0: The [N x C x ...] tensor of noiseless inputs.
+            t: The [N] tensor of diffusion steps [0-1].
+            noise: If specified, use this noise instead of generating new noise.
+
+        Returns:
+            x_t, the noisy version of x_0 under timestep t.
+        """
+        if noise is None:
+            noise = torch.randn_like(x_0)
+        assert noise.shape == x_0.shape, "noise must have same shape as x_0"
+
+        t = t.view(-1, *[1 for _ in range(len(x_0.shape) - 1)])
+        x_t = (1 - t) * x_0 + (self.sigma_min + (1 - self.sigma_min) * t) * noise
+
+        return x_t
+
+    def reverse_diffuse(self, x_t: torch.Tensor, t: torch.Tensor, noise: torch.Tensor) -> torch.Tensor:
+        """
+        Get original image from noisy version under timestep t.
+        """
+        assert noise.shape == x_t.shape, "noise must have same shape as x_t"
+        t = t.view(-1, *[1 for _ in range(len(x_t.shape) - 1)])
+        x_0 = (x_t - (self.sigma_min + (1 - self.sigma_min) * t) * noise) / (1 - t)
+        return x_0
+
+    def get_v(self, x_0: torch.Tensor, noise: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+        """
+        Compute the velocity of the diffusion process at time t.
+        """
+        return (1 - self.sigma_min) * noise - x_0
+
+    def get_cond(self, cond, **kwargs):
+        """
+        Get the conditioning data.
+        """
+        return cond
+    
+    def get_inference_cond(self, cond, **kwargs):
+        """
+        Get the conditioning data for inference.
+        """
+        return {'cond': cond, **kwargs}
+
+    def get_sampler(self, **kwargs) -> samplers.FlowEulerSampler:
+        """
+        Get the sampler for the diffusion process.
+        """
+        return samplers.FlowEulerSampler(self.sigma_min)
+    
+    def vis_cond(self, **kwargs):
+        """
+        Visualize the conditioning data.
+        """
+        return {}
+
+    def sample_t(self, batch_size: int) -> torch.Tensor:
+        """
+        Sample timesteps.
+        """
+        if self.t_schedule['name'] == 'uniform':
+            t = torch.rand(batch_size)
+        elif self.t_schedule['name'] == 'logitNormal':
+            mean = self.t_schedule['args']['mean']
+            std = self.t_schedule['args']['std']
+            t = torch.sigmoid(torch.randn(batch_size) * std + mean)
+        else:
+            raise ValueError(f"Unknown t_schedule: {self.t_schedule['name']}")
+        return t
+
+    def training_losses(
+        self,
+        x_0: torch.Tensor,
+        cond=None,
+        **kwargs
+    ) -> Tuple[Dict, Dict]:
+        """
+        Compute training losses for a single timestep.
+
+        Args:
+            x_0: The [N x C x ...] tensor of noiseless inputs.
+            cond: The [N x ...] tensor of additional conditions.
+            kwargs: Additional arguments to pass to the backbone.
+
+        Returns:
+            a dict with the key "loss" containing a tensor of shape [N].
+            may also contain other keys for different terms.
+        """
+        noise = torch.randn_like(x_0)
+        t = self.sample_t(x_0.shape[0]).to(x_0.device).float()
+        x_t = self.diffuse(x_0, t, noise=noise)
+        cond = self.get_cond(cond, **kwargs)
+        
+        pred = self.training_models['denoiser'](x_t, t * 1000, cond, **kwargs)
+        assert pred.shape == noise.shape == x_0.shape
+        target = self.get_v(x_0, noise, t)
+        terms = edict()
+        terms["mse"] = F.mse_loss(pred, target)
+        terms["loss"] = terms["mse"]
+
+        # log loss with time bins
+        mse_per_instance = np.array([
+            F.mse_loss(pred[i], target[i]).item()
+            for i in range(x_0.shape[0])
+        ])
+        time_bin = np.digitize(t.cpu().numpy(), np.linspace(0, 1, 11)) - 1
+        for i in range(10):
+            if (time_bin == i).sum() != 0:
+                terms[f"bin_{i}"] = {"mse": mse_per_instance[time_bin == i].mean()}
+
+        return terms, {}
+    
+    @torch.no_grad()
+    def run_snapshot(
+        self,
+        num_samples: int,
+        batch_size: int,
+        verbose: bool = False,
+    ) -> Dict:
+        # Use current step as seed to ensure different samples for each snapshot
+        import random
+        snapshot_seed = self.step
+        random.seed(snapshot_seed)
+        np.random.seed(snapshot_seed)
+        
+        g = torch.Generator()
+        g.manual_seed(snapshot_seed)
+        
+        dataloader = DataLoader(
+            copy.deepcopy(self.dataset),
+            batch_size=batch_size,
+            shuffle=True,
+            num_workers=0,
+            collate_fn=self.dataset.collate_fn if hasattr(self.dataset, 'collate_fn') else None,
+            generator=g,
+        )
+
+        # inference
+        sampler = self.get_sampler()
+        sample_gt = []
+        sample = []
+        cond_vis = []
+        sample_metadata = []
+        for i in range(0, num_samples, batch_size):
+            batch = min(batch_size, num_samples - i)
+            data = next(iter(dataloader))
+            data = {k: v[:batch].cuda() if isinstance(v, torch.Tensor) else v[:batch] for k, v in data.items()}
+            
+            # Collect metadata (dataset_name and sha256) for wandb display
+            if '_dataset_name' in data and '_sha256' in data:
+                for j in range(batch):
+                    sample_metadata.append(f"{data['_dataset_name'][j]}/{data['_sha256'][j]}")
+            
+            # Remove metadata fields before inference
+            data.pop('_dataset_name', None)
+            data.pop('_sha256', None)
+            
+            noise = torch.randn_like(data['x_0'])
+            sample_gt.append(data['x_0'])
+            cond_vis.append(self.vis_cond(**data))
+            del data['x_0']
+            args = self.get_inference_cond(**data)
+            res = sampler.sample(
+                self.models['denoiser'],
+                noise=noise,
+                **args,
+                steps=50, guidance_strength=3.0, verbose=verbose,
+            )
+            sample.append(res.samples)
+
+        sample_gt = torch.cat(sample_gt, dim=0)
+        sample = torch.cat(sample, dim=0)
+        sample_dict = {
+            'sample_gt': {'value': sample_gt, 'type': 'sample'},
+            'sample': {'value': sample, 'type': 'sample'},
+        }
+        if sample_metadata:
+            sample_dict['_metadata'] = sample_metadata
+        sample_dict.update(dict_reduce(cond_vis, None, {
+            'value': lambda x: torch.cat(x, dim=0),
+            'type': lambda x: x[0],
+        }))
+        
+        return sample_dict
+
+    
+class FlowMatchingCFGTrainer(ClassifierFreeGuidanceMixin, FlowMatchingTrainer):
+    """
+    Trainer for diffusion model with flow matching objective and classifier-free guidance.
+    
+    Args:
+        models (dict[str, nn.Module]): Models to train.
+        dataset (torch.utils.data.Dataset): Dataset.
+        output_dir (str): Output directory.
+        load_dir (str): Load directory.
+        step (int): Step to load.
+        batch_size (int): Batch size.
+        batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
+        batch_split (int): Split batch with gradient accumulation.
+        max_steps (int): Max steps.
+        optimizer (dict): Optimizer config.
+        lr_scheduler (dict): Learning rate scheduler config.
+        elastic (dict): Elastic memory management config.
+        grad_clip (float or dict): Gradient clip config.
+        ema_rate (float or list): Exponential moving average rates.
+        fp16_mode (str): FP16 mode.
+            - None: No FP16.
+            - 'inflat_all': Hold a inflated fp32 master param for all params.
+            - 'amp': Automatic mixed precision.
+        fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
+        finetune_ckpt (dict): Finetune checkpoint.
+        log_param_stats (bool): Log parameter stats.
+        i_print (int): Print interval.
+        i_log (int): Log interval.
+        i_sample (int): Sample interval.
+        i_save (int): Save interval.
+        i_ddpcheck (int): DDP check interval.
+
+        t_schedule (dict): Time schedule for flow matching.
+        sigma_min (float): Minimum noise level.
+        p_uncond (float): Probability of dropping conditions.
+    """
+    pass
+
+
+class TextConditionedFlowMatchingCFGTrainer(TextConditionedMixin, FlowMatchingCFGTrainer):
+    """
+    Trainer for text-conditioned diffusion model with flow matching objective and classifier-free guidance.
+    
+    Args:
+        models (dict[str, nn.Module]): Models to train.
+        dataset (torch.utils.data.Dataset): Dataset.
+        output_dir (str): Output directory.
+        load_dir (str): Load directory.
+        step (int): Step to load.
+        batch_size (int): Batch size.
+        batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
+        batch_split (int): Split batch with gradient accumulation.
+        max_steps (int): Max steps.
+        optimizer (dict): Optimizer config.
+        lr_scheduler (dict): Learning rate scheduler config.
+        elastic (dict): Elastic memory management config.
+        grad_clip (float or dict): Gradient clip config.
+        ema_rate (float or list): Exponential moving average rates.
+        fp16_mode (str): FP16 mode.
+            - None: No FP16.
+            - 'inflat_all': Hold a inflated fp32 master param for all params.
+            - 'amp': Automatic mixed precision.
+        fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
+        finetune_ckpt (dict): Finetune checkpoint.
+        log_param_stats (bool): Log parameter stats.
+        i_print (int): Print interval.
+        i_log (int): Log interval.
+        i_sample (int): Sample interval.
+        i_save (int): Save interval.
+        i_ddpcheck (int): DDP check interval.
+
+        t_schedule (dict): Time schedule for flow matching.
+        sigma_min (float): Minimum noise level.
+        p_uncond (float): Probability of dropping conditions.
+        text_cond_model(str): Text conditioning model.
+    """
+    pass
+
+
+class ImageConditionedFlowMatchingCFGTrainer(ImageConditionedMixin, FlowMatchingCFGTrainer):
+    """
+    Trainer for image-conditioned diffusion model with flow matching objective and classifier-free guidance.
+    
+    Args:
+        models (dict[str, nn.Module]): Models to train.
+        dataset (torch.utils.data.Dataset): Dataset.
+        output_dir (str): Output directory.
+        load_dir (str): Load directory.
+        step (int): Step to load.
+        batch_size (int): Batch size.
+        batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
+        batch_split (int): Split batch with gradient accumulation.
+        max_steps (int): Max steps.
+        optimizer (dict): Optimizer config.
+        lr_scheduler (dict): Learning rate scheduler config.
+        elastic (dict): Elastic memory management config.
+        grad_clip (float or dict): Gradient clip config.
+        ema_rate (float or list): Exponential moving average rates.
+        fp16_mode (str): FP16 mode.
+            - None: No FP16.
+            - 'inflat_all': Hold a inflated fp32 master param for all params.
+            - 'amp': Automatic mixed precision.
+        fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
+        finetune_ckpt (dict): Finetune checkpoint.
+        log_param_stats (bool): Log parameter stats.
+        i_print (int): Print interval.
+        i_log (int): Log interval.
+        i_sample (int): Sample interval.
+        i_save (int): Save interval.
+        i_ddpcheck (int): DDP check interval.
+
+        t_schedule (dict): Time schedule for flow matching.
+        sigma_min (float): Minimum noise level.
+        p_uncond (float): Probability of dropping conditions.
+        image_cond_model (str): Image conditioning model.
+    """
+    pass
+
+
+class ImageConditionedProjFlowMatchingCFGTrainer(ImageConditionedProjMixin, FlowMatchingCFGTrainer):
+    """
+    Trainer for image-conditioned diffusion model with view-aligned projection.
+    
+    Uses ImageConditionedProjMixin for 3D-to-2D feature projection with camera parameters.
+    CFG dropout is handled by ClassifierFreeGuidanceMixin (via p_uncond parameter).
+    
+    Args:
+        t_schedule (dict): Time schedule for flow matching.
+        sigma_min (float): Minimum noise level.
+        p_uncond (float): Probability of dropping conditions.
+        image_cond_model (dict): Image conditioning model config (DinoV3ProjFeatureExtractor).
+        run_projection_test (bool): Whether to run projection visualization test before training.
+    """
+    
+    def __init__(self, *args, run_projection_test: bool = True, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.run_projection_test = run_projection_test
+    
+    def training_losses(
+        self,
+        x_0: torch.Tensor,
+        cond=None,
+        **kwargs
+    ) -> Tuple[Dict, Dict]:
+        """
+        Compute training losses for a single timestep.
+        
+        Overridden to avoid passing extra kwargs to the model.
+
+        Args:
+            x_0: The [N x C x ...] tensor of noiseless inputs.
+            cond: The [N x ...] tensor of additional conditions.
+            kwargs: Additional arguments (camera info, view_idx, etc.) for conditioning.
+
+        Returns:
+            a dict with the key "loss" containing a tensor of shape [N].
+            may also contain other keys for different terms.
+        """
+        noise = torch.randn_like(x_0)
+        t = self.sample_t(x_0.shape[0]).to(x_0.device).float()
+        x_t = self.diffuse(x_0, t, noise=noise)
+        cond = self.get_cond(cond, **kwargs)
+        
+        # Note: SparseStructureFlowModel.forward() only accepts (x, t, cond)
+        # Do not pass extra kwargs to avoid unexpected keyword argument errors
+        pred = self.training_models['denoiser'](x_t, t * 1000, cond)
+        assert pred.shape == noise.shape == x_0.shape
+        target = self.get_v(x_0, noise, t)
+        terms = edict()
+        terms["mse"] = F.mse_loss(pred, target)
+        terms["loss"] = terms["mse"]
+
+        # log loss with time bins
+        mse_per_instance = np.array([
+            F.mse_loss(pred[i], target[i]).item()
+            for i in range(x_0.shape[0])
+        ])
+        time_bin = np.digitize(t.cpu().numpy(), np.linspace(0, 1, 11)) - 1
+        for i in range(10):
+            if (time_bin == i).sum() != 0:
+                terms[f"bin_{i}"] = {"mse": mse_per_instance[time_bin == i].mean()}
+
+        return terms, {}
+    
+    @torch.no_grad()
+    def run_snapshot(
+        self,
+        num_samples: int,
+        batch_size: int,
+        verbose: bool = False,
+    ) -> Dict:
+        """
+        Run snapshot with camera parameters for GT view rendering.
+        
+        Overrides parent to include camera parameters in sample dicts for
+        visualizing the GT camera view alongside the standard 4-view rendering.
+        """
+        # Use current step as seed to ensure different samples for each snapshot
+        import random
+        snapshot_seed = self.step
+        random.seed(snapshot_seed)
+        np.random.seed(snapshot_seed)
+        
+        g = torch.Generator()
+        g.manual_seed(snapshot_seed)
+        
+        dataloader = DataLoader(
+            copy.deepcopy(self.dataset),
+            batch_size=batch_size,
+            shuffle=True,
+            num_workers=0,
+            collate_fn=self.dataset.collate_fn if hasattr(self.dataset, 'collate_fn') else None,
+            generator=g,
+        )
+
+        # inference
+        sampler = self.get_sampler()
+        sample_gt_list = []
+        sample_list = []
+        cond_vis = []
+        sample_metadata = []
+        
+        # Camera params for GT view rendering
+        camera_distances = []
+        camera_angles = []
+        mesh_scales = []
+        
+        for i in range(0, num_samples, batch_size):
+            batch = min(batch_size, num_samples - i)
+            data = next(iter(dataloader))
+            data = {k: v[:batch].cuda() if isinstance(v, torch.Tensor) else v[:batch] for k, v in data.items()}
+            
+            # Collect metadata (dataset_name and sha256) for wandb display
+            if '_dataset_name' in data and '_sha256' in data:
+                for j in range(batch):
+                    sample_metadata.append(f"{data['_dataset_name'][j]}/{data['_sha256'][j]}")
+            
+            # Remove metadata fields before inference
+            data.pop('_dataset_name', None)
+            data.pop('_sha256', None)
+            
+            noise = torch.randn_like(data['x_0'])
+            
+            # Save GT sample
+            sample_gt_list.append(data['x_0'])
+            cond_vis.append(self.vis_cond(**data))
+            
+            # Save camera parameters for GT view rendering (if available)
+            if 'camera_distance' in data:
+                camera_distances.append(data['camera_distance'])
+            if 'camera_angle_x' in data:
+                camera_angles.append(data['camera_angle_x'])
+            if 'mesh_scale' in data:
+                mesh_scales.append(data['mesh_scale'])
+            
+            # Remove x_0 before inference
+            del data['x_0']
+            args = self.get_inference_cond(**data)
+            res = sampler.sample(
+                self.models['denoiser'],
+                noise=noise,
+                **args,
+                steps=50, guidance_strength=3.0, verbose=verbose,
+            )
+            sample_list.append(res.samples)
+
+        # Concatenate samples
+        sample_gt = torch.cat(sample_gt_list, dim=0)
+        sample = torch.cat(sample_list, dim=0)
+        
+        # Build sample dicts with camera info for GT view rendering
+        sample_gt_value = {'x_0': sample_gt}
+        sample_value = {'x_0': sample}
+        
+        # Add camera params if available
+        if len(camera_distances) > 0:
+            camera_distance = torch.cat(camera_distances, dim=0)
+            sample_gt_value['camera_distance'] = camera_distance
+            sample_value['camera_distance'] = camera_distance
+        if len(camera_angles) > 0:
+            camera_angle_x = torch.cat(camera_angles, dim=0)
+            sample_gt_value['camera_angle_x'] = camera_angle_x
+            sample_value['camera_angle_x'] = camera_angle_x
+        if len(mesh_scales) > 0:
+            mesh_scale = torch.cat(mesh_scales, dim=0)
+            sample_gt_value['mesh_scale'] = mesh_scale
+            sample_value['mesh_scale'] = mesh_scale
+        
+        sample_dict = {
+            'sample_gt': {'value': sample_gt_value, 'type': 'sample'},
+            'sample': {'value': sample_value, 'type': 'sample'},
+        }
+        if sample_metadata:
+            sample_dict['_metadata'] = sample_metadata
+        sample_dict.update(dict_reduce(cond_vis, None, {
+            'value': lambda x: torch.cat(x, dim=0),
+            'type': lambda x: x[0],
+        }))
+        
+        return sample_dict
+    
+    @torch.no_grad()
+    def visualize_sample(self, sample):
+        """
+        Convert a sample to images, including GT camera view if available.
+        
+        Args:
+            sample: Either a tensor or dict containing:
+                - 'x_0': latent tensor [B, C, D, H, W]
+                - 'camera_angle_x': [B] (optional)
+                - 'camera_distance': [B] (optional)
+                - 'mesh_scale': [B] (optional)
+                
+        Returns:
+            dict with visualization images or tensor
+        """
+        if hasattr(self.dataset, 'visualize_sample'):
+            if isinstance(sample, dict):
+                # Extract camera params if available
+                camera_angle_x = sample.get('camera_angle_x')
+                camera_distance = sample.get('camera_distance')
+                mesh_scale = sample.get('mesh_scale')
+                x_0 = sample.get('x_0', sample)
+                
+                return self.dataset.visualize_sample(
+                    x_0,
+                    camera_angle_x=camera_angle_x,
+                    camera_distance=camera_distance,
+                    mesh_scale=mesh_scale,
+                )
+            else:
+                return self.dataset.visualize_sample(sample)
+        else:
+            if isinstance(sample, dict):
+                return sample.get('x_0', sample)
+            return sample
+    
+    def run(self):
+        """
+        Run training with projection visualization test before starting.
+        """
+        # Run projection visualization test before training starts (if enabled)
+        if self.run_projection_test and self.is_master:
+            print('\n' + '='*60)
+            print('Running projection visualization test...')
+            print('='*60)
+            self._run_projection_visualization_test()
+            
+        super().run()
+    
+    @torch.no_grad()
+    def _run_projection_visualization_test(self, num_samples: int = 4):
+        """
+        Run projection visualization test on a few samples before training starts.
+        
+        This helps verify that the 3D-to-2D projection is working correctly.
+        """
+        import os
+        from torch.utils.data import DataLoader
+        
+        # Create a small dataloader
+        dataloader = DataLoader(
+            self.dataset,
+            batch_size=min(num_samples, self.snapshot_batch_size),
+            shuffle=True,
+            num_workers=0,
+            collate_fn=self.dataset.collate_fn if hasattr(self.dataset, 'collate_fn') else None,
+        )
+        
+        # Get one batch
+        data = next(iter(dataloader))
+        data = {k: v.cuda() if isinstance(v, torch.Tensor) else v for k, v in data.items()}
+        
+        # Extract condition image
+        cond = data.get('cond')
+        if cond is None:
+            print("Warning: No 'cond' field in data, skipping projection visualization test")
+            return
+        
+        # Save directory
+        save_dir = os.path.join(self.output_dir, 'samples', 'projection_test')
+        
+        # Call visualization method
+        if hasattr(self, 'visualize_projection_test'):
+            # Need to pass camera info as kwargs
+            kwargs = {k: v for k, v in data.items() if k != 'cond' and k != 'x_0'}
+            self.visualize_projection_test(
+                cond=cond,
+                save_dir=save_dir,
+                prefix="proj_test",
+                **kwargs
+            )
+            print(f"Projection visualization saved to: {save_dir}")
+        else:
+            print("Warning: visualize_projection_test not available")
diff --git a/pixal3d/trainers/flow_matching/mixins/classifier_free_guidance.py b/pixal3d/trainers/flow_matching/mixins/classifier_free_guidance.py
new file mode 100644
index 0000000..3ae3a4e
--- /dev/null
+++ b/pixal3d/trainers/flow_matching/mixins/classifier_free_guidance.py
@@ -0,0 +1,60 @@
+import torch
+import numpy as np
+from ....utils.general_utils import dict_foreach
+from ....pipelines import samplers
+
+
+class ClassifierFreeGuidanceMixin:
+    def __init__(self, *args, p_uncond: float = 0.1, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.p_uncond = p_uncond
+
+    def get_cond(self, cond, neg_cond=None, **kwargs):
+        """
+        Get the conditioning data.
+        """
+        assert neg_cond is not None, "neg_cond must be provided for classifier-free guidance" 
+
+        if self.p_uncond > 0:
+            # randomly drop the class label
+            def get_batch_size(cond):
+                if isinstance(cond, torch.Tensor):
+                    return cond.shape[0]
+                elif isinstance(cond, list):
+                    return len(cond)
+                else:
+                    raise ValueError(f"Unsupported type of cond: {type(cond)}")
+                
+            ref_cond = cond if not isinstance(cond, dict) else cond[list(cond.keys())[0]]
+            B = get_batch_size(ref_cond)
+            
+            def select(cond, neg_cond, mask):
+                if isinstance(cond, torch.Tensor):
+                    mask = torch.tensor(mask, device=cond.device).reshape(-1, *[1] * (cond.ndim - 1))
+                    return torch.where(mask, neg_cond, cond)
+                elif isinstance(cond, list):
+                    return [nc if m else c for c, nc, m in zip(cond, neg_cond, mask)]
+                else:
+                    raise ValueError(f"Unsupported type of cond: {type(cond)}")
+            
+            mask = list(np.random.rand(B) < self.p_uncond)
+            if not isinstance(cond, dict):
+                cond = select(cond, neg_cond, mask)
+            else:
+                # Apply select to each key in the dict
+                cond = {k: select(cond[k], neg_cond[k], mask) for k in cond.keys()}
+    
+        return cond
+
+    def get_inference_cond(self, cond, neg_cond=None, **kwargs):
+        """
+        Get the conditioning data for inference.
+        """
+        assert neg_cond is not None, "neg_cond must be provided for classifier-free guidance"
+        return {'cond': cond, 'neg_cond': neg_cond, **kwargs}
+    
+    def get_sampler(self, **kwargs) -> samplers.FlowEulerCfgSampler:
+        """
+        Get the sampler for the diffusion process.
+        """
+        return samplers.FlowEulerCfgSampler(self.sigma_min)
diff --git a/pixal3d/trainers/flow_matching/mixins/image_conditioned.py b/pixal3d/trainers/flow_matching/mixins/image_conditioned.py
new file mode 100644
index 0000000..13b932e
--- /dev/null
+++ b/pixal3d/trainers/flow_matching/mixins/image_conditioned.py
@@ -0,0 +1,249 @@
+from typing import *
+import torch
+import torch.nn.functional as F
+from torchvision import transforms
+from transformers import DINOv3ViTModel
+import numpy as np
+from PIL import Image
+
+from ....utils import dist_utils
+
+
+class DinoV2FeatureExtractor:
+    """
+    Feature extractor for DINOv2 models.
+    """
+    def __init__(self, model_name: str):
+        self.model_name = model_name
+        self.model = torch.hub.load('facebookresearch/dinov2', model_name, pretrained=True)
+        self.model.eval()
+        self.transform = transforms.Compose([
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ])
+
+    def to(self, device):
+        self.model.to(device)
+
+    def cuda(self):
+        self.model.cuda()
+
+    def cpu(self):
+        self.model.cpu()
+    
+    @torch.no_grad()
+    def __call__(self, image: Union[torch.Tensor, List[Image.Image]]) -> torch.Tensor:
+        """
+        Extract features from the image.
+        
+        Args:
+            image: A batch of images as a tensor of shape (B, C, H, W) or a list of PIL images.
+        
+        Returns:
+            A tensor of shape (B, N, D) where N is the number of patches and D is the feature dimension.
+        """
+        if isinstance(image, torch.Tensor):
+            assert image.ndim == 4, "Image tensor should be batched (B, C, H, W)"
+        elif isinstance(image, list):
+            assert all(isinstance(i, Image.Image) for i in image), "Image list should be list of PIL images"
+            image = [i.resize((518, 518), Image.LANCZOS) for i in image]
+            image = [np.array(i.convert('RGB')).astype(np.float32) / 255 for i in image]
+            image = [torch.from_numpy(i).permute(2, 0, 1).float() for i in image]
+            image = torch.stack(image).cuda()
+        else:
+            raise ValueError(f"Unsupported type of image: {type(image)}")
+        
+        image = self.transform(image).cuda()
+        features = self.model(image, is_training=True)['x_prenorm']
+        patchtokens = F.layer_norm(features, features.shape[-1:])
+        return patchtokens
+    
+
+class DinoV3FeatureExtractor:
+    """
+    Feature extractor for DINOv3 models.
+    """
+    def __init__(self, model_name: str, image_size=512):
+        self.model_name = model_name
+        self.model = DINOv3ViTModel.from_pretrained(model_name)
+        self.model.eval()
+        self.image_size = image_size
+        self.transform = transforms.Compose([
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ])
+
+    def to(self, device):
+        self.model.to(device)
+
+    def cuda(self):
+        self.model.cuda()
+
+    def cpu(self):
+        self.model.cpu()
+
+    def extract_features(self, image: torch.Tensor) -> torch.Tensor:
+        image = image.to(self.model.embeddings.patch_embeddings.weight.dtype)
+        hidden_states = self.model.embeddings(image, bool_masked_pos=None)
+        position_embeddings = self.model.rope_embeddings(image)
+
+        for i, layer_module in enumerate(self.model.layer):
+            hidden_states = layer_module(
+                hidden_states,
+                position_embeddings=position_embeddings,
+            )
+
+        return F.layer_norm(hidden_states, hidden_states.shape[-1:])
+        
+    @torch.no_grad()
+    def __call__(self, image: Union[torch.Tensor, List[Image.Image]]) -> torch.Tensor:
+        """
+        Extract features from the image.
+        
+        Args:
+            image: A batch of images as a tensor of shape (B, C, H, W) or a list of PIL images.
+        
+        Returns:
+            A tensor of shape (B, N, D) where N is the number of patches and D is the feature dimension.
+        """
+        if isinstance(image, torch.Tensor):
+            assert image.ndim == 4, "Image tensor should be batched (B, C, H, W)"
+        elif isinstance(image, list):
+            assert all(isinstance(i, Image.Image) for i in image), "Image list should be list of PIL images"
+            image = [i.resize((self.image_size, self.image_size), Image.LANCZOS) for i in image]
+            image = [np.array(i.convert('RGB')).astype(np.float32) / 255 for i in image]
+            image = [torch.from_numpy(i).permute(2, 0, 1).float() for i in image]
+            image = torch.stack(image).cuda()
+        else:
+            raise ValueError(f"Unsupported type of image: {type(image)}")
+        
+        image = self.transform(image).cuda()
+        features = self.extract_features(image)
+        return features
+    
+
+class ImageConditionedMixin:
+    """
+    Mixin for image-conditioned models.
+    
+    Args:
+        image_cond_model: The image conditioning model.
+    """
+    def __init__(self, *args, image_cond_model: dict, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.image_cond_model_config = image_cond_model
+        self.image_cond_model = None     # the model is init lazily
+        
+    def _init_image_cond_model(self):
+        """
+        Initialize the image conditioning model.
+        """
+        with dist_utils.local_master_first():
+            self.image_cond_model = globals()[self.image_cond_model_config['name']](**self.image_cond_model_config.get('args', {}))
+            self.image_cond_model.cuda()
+    
+    @torch.no_grad()
+    def encode_image(self, image: Union[torch.Tensor, List[Image.Image]]) -> torch.Tensor:
+        """
+        Encode the image.
+        """
+        if self.image_cond_model is None:
+            self._init_image_cond_model()
+        features = self.image_cond_model(image)
+        return features
+        
+    def get_cond(self, cond, **kwargs):
+        """
+        Get the conditioning data.
+        """
+        cond = self.encode_image(cond)
+        kwargs['neg_cond'] = torch.zeros_like(cond)
+        cond = super().get_cond(cond, **kwargs)
+        return cond
+    
+    def get_inference_cond(self, cond, **kwargs):
+        """
+        Get the conditioning data for inference.
+        """
+        cond = self.encode_image(cond)
+        kwargs['neg_cond'] = torch.zeros_like(cond)
+        cond = super().get_inference_cond(cond, **kwargs)
+        return cond
+
+    def vis_cond(self, cond, **kwargs):
+        """
+        Visualize the conditioning data.
+        """
+        return {'image': {'value': cond, 'type': 'image'}}
+    
+
+class MultiImageConditionedMixin:
+    """
+    Mixin for multiple-image-conditioned models.
+    
+    Args:
+        image_cond_model: The image conditioning model.
+    """
+    def __init__(self, *args, image_cond_model: dict, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.image_cond_model_config = image_cond_model
+        self.image_cond_model = None     # the model is init lazily
+        
+    def _init_image_cond_model(self):
+        """
+        Initialize the image conditioning model.
+        """
+        with dist_utils.local_master_first():
+            self.image_cond_model = globals()[self.image_cond_model_config['name']](**self.image_cond_model_config.get('args', {}))
+    
+    @torch.no_grad()
+    def encode_images(self, images: Union[List[torch.Tensor], List[List[Image.Image]]]) -> List[torch.Tensor]:
+        """
+        Encode the image.
+        """
+        if self.image_cond_model is None:
+            self._init_image_cond_model()
+        seqlen = [len(i) for i in images]
+        images = torch.cat(images, dim=0) if isinstance(images[0], torch.Tensor) else sum(images, [])
+        features = self.image_cond_model(images)
+        features = torch.split(features, seqlen)
+        features = [feature.reshape(-1, feature.shape[-1]) for feature in features]
+        return features
+        
+    def get_cond(self, cond, **kwargs):
+        """
+        Get the conditioning data.
+        """
+        cond = self.encode_images(cond)
+        kwargs['neg_cond'] = [
+            torch.zeros_like(cond[0][:1, :]) for _ in range(len(cond))
+        ]
+        cond = super().get_cond(cond, **kwargs)
+        return cond
+    
+    def get_inference_cond(self, cond, **kwargs):
+        """
+        Get the conditioning data for inference.
+        """
+        cond = self.encode_images(cond)
+        kwargs['neg_cond'] = [
+            torch.zeros_like(cond[0][:1, :]) for _ in range(len(cond))
+        ]
+        cond = super().get_inference_cond(cond, **kwargs)
+        return cond
+
+    def vis_cond(self, cond, **kwargs):
+        """
+        Visualize the conditioning data.
+        """
+        H, W = cond[0].shape[-2:]
+        vis = []
+        for images in cond:
+            canvas = torch.zeros(3, H * 2, W * 2, device=images.device, dtype=images.dtype)
+            for i, image in enumerate(images):
+                if i == 4:
+                    break
+                kh = i // 2
+                kw = i % 2
+                canvas[:, kh*H:(kh+1)*H, kw*W:(kw+1)*W] = image
+            vis.append(canvas)
+        vis = torch.stack(vis)
+        return {'image': {'value': vis, 'type': 'image'}}
diff --git a/pixal3d/trainers/flow_matching/mixins/image_conditioned_proj.py b/pixal3d/trainers/flow_matching/mixins/image_conditioned_proj.py
new file mode 100644
index 0000000..ccafe87
--- /dev/null
+++ b/pixal3d/trainers/flow_matching/mixins/image_conditioned_proj.py
@@ -0,0 +1,1530 @@
+"""
+View-Aligned (Projection) Image Conditioned Mixin for Pixal3D
+
+This module implements DINOv3-based feature extraction with view-aligned projection,
+supporting camera-aware 3D-to-2D feature mapping.
+"""
+
+from typing import *
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import transforms
+from transformers import DINOv3ViTModel
+import numpy as np
+from PIL import Image, ImageDraw
+
+import torch.distributed as dist
+from ....utils import dist_utils
+from ....utils.dist_utils import read_file_dist
+
+
+# =============================================================================
+# Projection Utilities
+# =============================================================================
+
+def project_points_to_image_batch(
+    points_3d: torch.Tensor, 
+    transform_matrix: torch.Tensor, 
+    camera_angle_x: torch.Tensor, 
+    resolution: int = 518
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Project 3D points to 2D image coordinates (batch processing).
+    
+    Args:
+        points_3d: torch.Tensor, shape [N, 3] or [B, N, 3], 3D point coordinates (in [-1, 1] range)
+        transform_matrix: torch.Tensor, shape [B, 4, 4], camera transformation matrix
+        camera_angle_x: torch.Tensor, shape [B], horizontal field of view angle (radians)
+        resolution: int, image resolution, default 518
+    
+    Returns:
+        points_2d: torch.Tensor, shape [B, N, 2], image coordinates [x, y]
+        depth: torch.Tensor, shape [B, N], depth values
+        valid_mask: torch.Tensor, shape [B, N], mask for points within view
+    """
+    device = points_3d.device
+    B = transform_matrix.shape[0]
+    
+    # Ensure inputs are torch.Tensor on correct device
+    if not isinstance(transform_matrix, torch.Tensor):
+        transform_matrix = torch.tensor(transform_matrix, dtype=torch.float32, device=device)
+    if not isinstance(points_3d, torch.Tensor):
+        points_3d = torch.tensor(points_3d, dtype=torch.float32, device=device)
+    if not isinstance(camera_angle_x, torch.Tensor):
+        camera_angle_x = torch.tensor(camera_angle_x, dtype=torch.float32, device=device)
+    
+    # Expand points_3d to batch dimension: [N, 3] -> [B, N, 3]
+    if points_3d.dim() == 2:
+        points_3d_batch = points_3d.unsqueeze(0).expand(B, -1, -1)
+    else:
+        points_3d_batch = points_3d
+    N = points_3d_batch.shape[1]
+    
+    # Add homogeneous coordinates: [B, N, 3] -> [B, N, 4]
+    ones = torch.ones(B, N, 1, device=device, dtype=points_3d_batch.dtype)
+    points_homogeneous = torch.cat([points_3d_batch, ones], dim=-1)  # [B, N, 4]
+    
+    # Compute world to camera transformation matrix
+    world_to_camera = torch.linalg.inv(transform_matrix.float()).to(transform_matrix.dtype)  # linalg.inv requires fp32+
+    
+    # Batch transform to camera coordinate system: [B, N, 4] @ [B, 4, 4]^T -> [B, N, 3]
+    points_camera = torch.bmm(points_homogeneous, world_to_camera.transpose(-2, -1))[..., :3]  # [B, N, 3]
+    
+    # Extract camera coordinates
+    x_cam = points_camera[..., 0]  # [B, N]
+    y_cam = points_camera[..., 1]  # [B, N]
+    z_cam = points_camera[..., 2]  # [B, N]
+    
+    # Depth value (Z value in camera coordinate system, note Blender camera faces -Z direction)
+    depth = -z_cam  # [B, N]
+    
+    # Compute camera intrinsics (batch processing)
+    sensor_width = 32.0  # mm
+    focal_length = 16.0 / torch.tan(camera_angle_x / 2.0)  # [B]
+    focal_length_pixels = focal_length * resolution / sensor_width  # [B]
+    
+    # Expand focal_length_pixels dimension for broadcasting: [B] -> [B, 1]
+    focal_length_pixels = focal_length_pixels.unsqueeze(1)  # [B, 1]
+    
+    # Perspective projection to NDC coordinates
+    x_ndc = focal_length_pixels * x_cam / (-z_cam + 1e-8)  # [B, N]
+    y_ndc = focal_length_pixels * y_cam / (-z_cam + 1e-8)  # [B, N]
+    
+    # Convert to image coordinates (pixel coordinates)
+    x_pixel = x_ndc + resolution / 2.0  # [B, N]
+    y_pixel = -y_ndc + resolution / 2.0  # [B, N], flip Y axis
+    
+    # Create validity mask (points within image range and in front of camera)
+    valid_mask = (
+        (x_pixel >= 0) & (x_pixel < resolution) &
+        (y_pixel >= 0) & (y_pixel < resolution) &
+        (depth > 0)  # In front of camera
+    )  # [B, N]
+    
+    points_2d = torch.stack([x_pixel, y_pixel], dim=-1)  # [B, N, 2]
+    
+    return points_2d, depth, valid_mask
+
+
+def sample_features(fmap: torch.Tensor, queries_ndc: torch.Tensor) -> torch.Tensor:
+    """
+    Sample features from feature map at specified NDC coordinates.
+    
+    Args:
+        fmap: torch.Tensor, shape [B, C, H, W], feature map
+        queries_ndc: torch.Tensor, shape [B, K, 2], normalized device coordinates
+    
+    Returns:
+        torch.Tensor, shape [B, C, K], sampled features
+    """
+    B, C, H, W = fmap.shape
+    Bq, K, _ = queries_ndc.shape
+    assert Bq == B, "Batch size mismatch"
+
+    # grid_sample requires (B, out_h, out_w, 2), here we want K points -> out_h=K, out_w=1
+    grid = queries_ndc.view(B, K, 1, 2)  # (B, K, 1, 2)
+
+    # Bilinear interpolation, align_corners=False (consistent with [-1,1] pixel center convention)
+    feat = F.grid_sample(
+        fmap, grid, mode='bilinear',
+        align_corners=False, padding_mode='border'  # border avoids out-of-bound becoming 0
+    )  # (B, C, K, 1)
+
+    return feat.squeeze(-1)  # (B, C, K)
+
+
+# =============================================================================
+# Projection Grid Module
+# =============================================================================
+
+class ProjGrid(nn.Module):
+    """
+    3D Grid Projection Module.
+    
+    Projects a 3D grid of points to 2D image coordinates and samples features 
+    from the image feature map at those locations.
+    
+    This is the core module for view-aligned feature extraction.
+    """
+    def __init__(self, grid_resolution: int = 16, image_resolution: int = 518):
+        super().__init__()
+        self.grid_resolution = grid_resolution
+        self.image_resolution = image_resolution
+        
+        # Create 3D grid points
+        one_dim = torch.linspace(-1, 1, grid_resolution)
+        x, y, z = torch.meshgrid(one_dim, one_dim, one_dim, indexing='ij')
+        grid_points = torch.stack((x, y, z), dim=-1)
+        
+        # Rotation matrix to align with Blender coordinate system
+        rotation_matrix = torch.tensor([
+            [1.0, 0.0, 0.0],
+            [0.0, 0.0, -1.0],
+            [0.0, 1.0, 0.0]
+        ])
+        grid_points = torch.matmul(grid_points, rotation_matrix.T)
+        grid_points = grid_points.reshape(-1, 3)
+        self.register_buffer('grid_points', grid_points)  # [R³, 3]
+        
+        # Default front view transformation matrix
+        front_view_transform_matrix = torch.tensor([
+            [1.0, 0.0, 0.0, 0.0],
+            [0.0, 0.0, -1.0, -2.0],
+            [0.0, 1.0, 0.0, 0.0],
+            [0.0, 0.0, 0.0, 1.0]
+        ])
+        self.register_buffer("front_view_transform_matrix", front_view_transform_matrix)
+        
+    def forward(
+        self, 
+        features_map: torch.Tensor, 
+        camera_angle_x: torch.Tensor,
+        distance: torch.Tensor,
+        mesh_scale: torch.Tensor,
+        transform_matrix: Optional[torch.Tensor] = None,
+        BHWC: bool = True
+    ) -> torch.Tensor:
+        """
+        Project 3D grid points to image and sample features.
+        
+        Args:
+            features_map: Feature map, shape [B, H, W, C] if BHWC else [B, C, H, W]
+            camera_angle_x: Camera FOV angle, shape [B]
+            distance: Camera distance, shape [B]
+            mesh_scale: Mesh scale factor, shape [B]
+            transform_matrix: Optional camera transform matrix, shape [B, 4, 4]
+            BHWC: Whether features_map is in BHWC format
+            
+        Returns:
+            Projected features, shape [B, grid_resolution³, C]
+        """
+        if BHWC:
+            B, H, W, C = features_map.shape
+        else:
+            B, C, H, W = features_map.shape
+            
+        grid_points = self.grid_points
+        grid_points = grid_points.expand(B, -1, -1)
+        grid_points = grid_points / mesh_scale.unsqueeze(-1).unsqueeze(-1) / 2  # Scale alignment
+        assert transform_matrix is None, "transform_matrix is not None"
+        if transform_matrix is None:
+            transform_matrix = self.front_view_transform_matrix
+            transform_matrix = transform_matrix.expand(B, -1, -1).clone()
+            transform_matrix[:, 1, 3] = -distance  # Set camera distance
+            
+        # Project to image coordinates (simulate Blender projection)
+        image_points, depth, valid_mask = project_points_to_image_batch(
+            grid_points, transform_matrix, camera_angle_x, self.image_resolution
+        )
+        
+        # Normalize to [-1, 1] for grid_sample
+        image_points_norm = (image_points + 0.5) / self.image_resolution * 2 - 1
+        
+        if BHWC:
+            features_map = features_map.permute(0, 3, 1, 2)  # [B, C, H, W]
+            
+        # Sample features from DINOv3 patch feature map
+        x = sample_features(features_map, image_points_norm)  # [B, C, K]
+        x = x.permute(0, 2, 1)  # [B, K, C]
+   
+        return x
+    
+    def visualize_projection(
+        self,
+        image: torch.Tensor,
+        camera_angle_x: torch.Tensor,
+        distance: torch.Tensor,
+        mesh_scale: torch.Tensor,
+        transform_matrix: Optional[torch.Tensor] = None,
+        save_dir: Optional[str] = None,
+        prefix: str = "proj_vis",
+    ) -> List[Image.Image]:
+        """
+        Visualize the projected 3D grid points on the input image.
+        
+        Args:
+            image: Input image tensor [B, C, H, W], assumed to be in [0, 1] range
+            camera_angle_x: Camera FOV angle, shape [B]
+            distance: Camera distance, shape [B]
+            mesh_scale: Mesh scale factor, shape [B]
+            transform_matrix: Optional camera transform matrix, shape [B, 4, 4]
+            save_dir: Directory to save visualizations (optional)
+            prefix: Prefix for saved files
+            
+        Returns:
+            List of PIL Images with projected points overlaid
+        """
+        B = image.shape[0]
+        
+        # Get projected points
+        grid_points = self.grid_points.expand(B, -1, -1)
+        grid_points = grid_points / mesh_scale.unsqueeze(-1).unsqueeze(-1) / 2
+        assert transform_matrix is None, "transform_matrix is not None"
+        if transform_matrix is None:
+            transform_matrix = self.front_view_transform_matrix
+            transform_matrix = transform_matrix.expand(B, -1, -1).clone()
+            transform_matrix[:, 1, 3] = -distance
+            
+        image_points, depth, valid_mask = project_points_to_image_batch(
+            grid_points, transform_matrix, camera_angle_x, self.image_resolution
+        )
+        
+        # Convert image to PIL for visualization
+        vis_images = []
+        for b in range(B):
+            # Convert tensor to PIL image
+            img_np = image[b].cpu().permute(1, 2, 0).numpy()
+            img_np = (img_np * 255).clip(0, 255).astype(np.uint8)
+            
+            # Resize to image_resolution if needed
+            pil_img = Image.fromarray(img_np)
+            if pil_img.size != (self.image_resolution, self.image_resolution):
+                pil_img = pil_img.resize((self.image_resolution, self.image_resolution), Image.LANCZOS)
+            
+            # Create a copy for drawing
+            vis_img = pil_img.copy()
+            draw = ImageDraw.Draw(vis_img)
+            
+            # Get points for this batch
+            pts = image_points[b].cpu().numpy()  # [K, 2]
+            depths = depth[b].cpu().numpy()  # [K]
+            mask = valid_mask[b].cpu().numpy()  # [K]
+            
+            # Normalize depth for coloring
+            valid_depths = depths[mask]
+            if len(valid_depths) > 0:
+                d_min, d_max = valid_depths.min(), valid_depths.max()
+                if d_max - d_min > 1e-6:
+                    depths_norm = (depths - d_min) / (d_max - d_min)
+                else:
+                    depths_norm = np.ones_like(depths) * 0.5
+            else:
+                depths_norm = np.ones_like(depths) * 0.5
+            
+            # Draw projected points
+            R = self.grid_resolution
+            for i, (pt, d, m, dn) in enumerate(zip(pts, depths, mask, depths_norm)):
+                if not m:
+                    continue
+                    
+                x, y = pt
+                
+                # Color by depth (blue=near, red=far)
+                r = int(255 * dn)
+                g = int(255 * (1 - abs(2 * dn - 1)))
+                b_color = int(255 * (1 - dn))
+                color = (r, g, b_color)
+                
+                # Draw small circle
+                radius = 2
+                draw.ellipse(
+                    [x - radius, y - radius, x + radius, y + radius],
+                    fill=color,
+                    outline=color
+                )
+            
+            vis_images.append(vis_img)
+            
+            # Save if directory is specified
+            if save_dir is not None:
+                os.makedirs(save_dir, exist_ok=True)
+                save_path = os.path.join(save_dir, f"{prefix}_batch{b}.png")
+                vis_img.save(save_path)
+                print(f"Saved projection visualization to: {save_path}")
+        
+        return vis_images
+
+
+# =============================================================================
+# DINOv3 Feature Extractor with Projection
+# =============================================================================
+
+class DinoV3ProjFeatureExtractor(nn.Module):
+    """
+    DINOv3 Feature Extractor with View-Aligned Projection.
+    
+    This extractor produces both:
+    1. Global features (CLS token + register tokens) in embed_dim
+    2. View-aligned projected features (3D grid projected to 2D and sampled)
+       - Without NAF: [B, R³, embed_dim]
+       - With NAF:    [B, R³, embed_dim * 2]  (concat of lr and hr features)
+    
+    NOTE: proj_linear has been moved to per-block ProjectAttention / SparseProjectAttention.
+    This module now outputs raw DINOv3 features for proj (optionally concatenated with NAF-upsampled features).
+    
+    Args:
+        model_name: Name of the pretrained DINOv3 model
+        image_size: Input image size (default: 512)
+        grid_resolution: Resolution of the 3D projection grid (default: 16)
+        use_naf_upsample: Whether to use NAF to upsample features (default: False)
+        naf_target_size: Target spatial size for NAF upsampling (default: [128, 128])
+    """
+    def __init__(
+        self, 
+        model_name: str,
+        image_size: int = 512,
+        grid_resolution: int = 16,
+        use_naf_upsample: bool = False,
+        naf_target_size: Optional[List[int]] = None,
+    ):
+        super().__init__()
+        self.model_name = model_name
+        self.image_size = image_size
+        self.grid_resolution = grid_resolution
+        self.use_naf_upsample = use_naf_upsample
+        if naf_target_size is None:
+            self.naf_target_size = (128, 128)
+        elif isinstance(naf_target_size, int):
+            self.naf_target_size = (naf_target_size, naf_target_size)
+        else:
+            self.naf_target_size = tuple(naf_target_size)
+        
+        # Load DINOv3 model (frozen, no trainable params in this module)
+        self.model = DINOv3ViTModel.from_pretrained(model_name)
+        self.model.eval()
+        self.model.requires_grad_(False)
+        
+        # Image transform (only normalize, no resize - assume already resized)
+        self.transform = transforms.Compose([
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ])
+        
+        # Get patch info
+        self.patch_size = self.model.config.patch_size
+        self.patch_number = image_size // self.patch_size
+        self.embed_dim = self.model.config.hidden_size
+        
+        # Projection grid for view-aligned features
+        self.proj_grid = ProjGrid(
+            grid_resolution=grid_resolution,
+            image_resolution=image_size
+        )
+        
+        # NAF upsampler (frozen, no trainable params)
+        self.naf_model = None  # Lazy-loaded on first use to avoid import if not needed
+        
+        # proj_channels: the output dimension of proj features
+        # Without NAF: embed_dim (e.g. 1024)
+        # With NAF: embed_dim * 2 (e.g. 2048, concat of lr and hr)
+        self.proj_channels = self.embed_dim * 2 if use_naf_upsample else self.embed_dim
+        
+        # NOTE: proj_linear removed — now lives in each denoiser block's ProjectAttention
+    
+    def _load_naf(self):
+        """Lazy-load pretrained NAF model."""
+        if self.naf_model is None:
+            import torch.hub
+            device = next(self.model.parameters()).device
+            self.naf_model = torch.hub.load(
+                "valeoai/NAF", "naf", pretrained=True, device=device, trust_repo=True
+            )
+            self.naf_model.eval()
+            self.naf_model.requires_grad_(False)
+        
+    def to(self, device):
+        super().to(device)
+        self.model.to(device)
+        self.proj_grid.to(device)
+        if self.naf_model is not None:
+            self.naf_model.to(device)
+        return self
+
+    def cuda(self):
+        super().cuda()
+        self.model.cuda()
+        self.proj_grid.cuda()
+        if self.naf_model is not None:
+            self.naf_model.cuda()
+        return self
+
+    def cpu(self):
+        super().cpu()
+        self.model.cpu()
+        self.proj_grid.cpu()
+        if self.naf_model is not None:
+            self.naf_model.cpu()
+        return self
+    
+    def extract_features(self, image: torch.Tensor) -> torch.Tensor:
+        """Extract features using DINOv3."""
+        image = image.to(self.model.embeddings.patch_embeddings.weight.dtype)
+        hidden_states = self.model.embeddings(image, bool_masked_pos=None)
+        position_embeddings = self.model.rope_embeddings(image)
+
+        for layer_module in self.model.layer:
+            hidden_states = layer_module(
+                hidden_states,
+                position_embeddings=position_embeddings,
+            )
+
+        return F.layer_norm(hidden_states, hidden_states.shape[-1:])
+    
+    def forward(
+        self,
+        image: Union[torch.Tensor, List[Image.Image]],
+        camera_angle_x: Optional[torch.Tensor] = None,
+        distance: Optional[torch.Tensor] = None,
+        mesh_scale: Optional[torch.Tensor] = None,
+        transform_matrix: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Extract view-aligned features from the image.
+        
+        Args:
+            image: Input image tensor [B, C, H, W] or list of PIL images
+            camera_angle_x: Camera FOV angle in radians [B]
+            distance: Camera distance [B]
+            mesh_scale: Mesh scale factor [B]
+            transform_matrix: Optional camera transform matrix [B, 4, 4]
+        
+        Returns:
+            Tuple of (global_features, proj_features):
+            - global_features: [B, num_global_tokens, embed_dim]
+            - proj_features: [B, grid_resolution³, proj_channels]
+              where proj_channels = embed_dim (no NAF) or embed_dim*2 (with NAF)
+        """
+        # Handle input types
+        if isinstance(image, torch.Tensor):
+            assert image.ndim == 4, "Image tensor should be batched (B, C, H, W)"
+        elif isinstance(image, list):
+            assert all(isinstance(i, Image.Image) for i in image), "Image list should be list of PIL images"
+            image = [i.resize((self.image_size, self.image_size), Image.LANCZOS) for i in image]
+            image = [np.array(i.convert('RGB')).astype(np.float32) / 255 for i in image]
+            image = [torch.from_numpy(i).permute(2, 0, 1).float() for i in image]
+            image = torch.stack(image).cuda()
+        else:
+            raise ValueError(f"Unsupported type of image: {type(image)}")
+        
+        B = image.shape[0]
+        
+        # Keep a copy of the unnormalized image for NAF guide
+        if self.use_naf_upsample:
+            image_for_naf = image.clone()  # [B, 3, H, W], in [0, 1] range
+        
+        # Apply transform (ImageNet normalization)
+        image = self.transform(image)
+        
+        # Extract DINOv3 features (frozen, no gradients)
+        with torch.no_grad():
+            z = self.extract_features(image)
+            
+            # Split into CLS token, register tokens, and patch tokens
+            z_clstoken = z[:, 0:1]  # [B, 1, D]
+            num_reg = getattr(self.model.config, 'num_register_tokens', 4)
+            z_regtokens = z[:, 1:1+num_reg]  # [B, num_reg, D]
+            z_patchtokens = z[:, 1+num_reg:]  # [B, num_patches, D]
+            
+            # Reshape patch tokens to spatial grid: [B, h, w, D]
+            z_patchtokens_spatial = z_patchtokens.reshape(
+                B, self.patch_number, self.patch_number, -1
+            )  # [B, h, w, D]
+            
+            if camera_angle_x is None or distance is None or mesh_scale is None:
+                raise ValueError("camera_angle_x, distance, and mesh_scale must be provided")
+            
+            # --- Low-resolution branch: sample from DINOv3 patch feature map ---
+            z_proj_lr = self.proj_grid(
+                z_patchtokens_spatial, 
+                camera_angle_x, 
+                distance, 
+                mesh_scale,
+                transform_matrix
+            )  # [B, grid_res³, D]
+            
+            # --- High-resolution branch (NAF): upsample then sample ---
+            if self.use_naf_upsample:
+                self._load_naf()
+                # NAF expects: guide [B, 3, H, W], lr_features [B, C, h, w], target_size (H', W')
+                lr_features_bchw = z_patchtokens_spatial.permute(0, 3, 1, 2)  # [B, D, h, w]
+                hr_features = self.naf_model(
+                    image_for_naf, lr_features_bchw, self.naf_target_size
+                )  # [B, D, H', W']
+                
+                # Sample from high-res feature map using same projection coordinates
+                z_proj_hr = self.proj_grid(
+                    hr_features,
+                    camera_angle_x,
+                    distance,
+                    mesh_scale,
+                    transform_matrix,
+                    BHWC=False  # hr_features is [B, C, H', W']
+                )  # [B, grid_res³, D]
+                
+                # Concatenate lr and hr: [B, grid_res³, D*2]
+                z_proj = torch.cat([z_proj_lr, z_proj_hr], dim=-1)
+            else:
+                z_proj = z_proj_lr  # [B, grid_res³, D]
+                
+            # Combine global tokens
+            z_global = torch.cat([z_clstoken, z_regtokens], dim=1)  # [B, 1+num_reg, D]
+        
+        # proj_linear has been moved to per-block ProjectAttention
+        # z_proj stays in proj_channels, each block will project independently
+        
+        return z_global, z_proj
+    
+    @torch.no_grad()
+    def visualize_projection(
+        self,
+        image: torch.Tensor,
+        camera_angle_x: torch.Tensor,
+        distance: torch.Tensor,
+        mesh_scale: torch.Tensor,
+        transform_matrix: Optional[torch.Tensor] = None,
+        save_dir: Optional[str] = None,
+        prefix: str = "proj_vis",
+    ) -> List[Image.Image]:
+        """
+        Visualize the projected 3D grid points on the input image.
+        
+        This is a convenience method that delegates to ProjGrid.visualize_projection.
+        
+        Args:
+            image: Input image tensor [B, C, H, W], in [0, 1] range (before ImageNet normalization)
+            camera_angle_x: Camera FOV angle, shape [B]
+            distance: Camera distance, shape [B]
+            mesh_scale: Mesh scale factor, shape [B]
+            transform_matrix: Optional camera transform matrix, shape [B, 4, 4]
+            save_dir: Directory to save visualizations (optional)
+            prefix: Prefix for saved files
+            
+        Returns:
+            List of PIL Images with projected points overlaid
+        """
+        return self.proj_grid.visualize_projection(
+            image=image,
+            camera_angle_x=camera_angle_x,
+            distance=distance,
+            mesh_scale=mesh_scale,
+            transform_matrix=transform_matrix,
+            save_dir=save_dir,
+            prefix=prefix,
+        )
+
+
+# =============================================================================
+# DINOv3 + VAE Gated Feature Extractor with Projection
+# =============================================================================
+
+class DinoV3VaeProjFeatureExtractor(nn.Module):
+    """
+    DINOv3 + Flux VAE Feature Extractor with Gated Fusion and View-Aligned Projection.
+    
+    Produces three outputs for GatedProjectAttention:
+    1. Global features (CLS + register tokens from DINOv3) for cross-attention
+    2. Semantic proj features (DINOv3 patch tokens projected to 3D grid)
+    3. Color proj features (Flux VAE latent projected to 3D grid)
+    
+    Both DINOv3 and VAE are frozen. The gated fusion happens inside each
+    denoiser block's GatedProjectAttention module (trainable gate + proj_linears).
+    
+    Args:
+        dino_model_name: Pretrained DINOv3 model name
+        vae_model_name: Pretrained Flux VAE model name
+        image_size: Input image size (default: 512)
+        grid_resolution: Resolution of the 3D projection grid (default: 16)
+    """
+    def __init__(
+        self,
+        dino_model_name: str,
+        vae_model_name: str = "black-forest-labs/FLUX.1-dev",
+        image_size: int = 512,
+        grid_resolution: int = 16,
+    ):
+        super().__init__()
+        self.image_size = image_size
+        self.grid_resolution = grid_resolution
+        
+        # --- DINOv3 backbone (frozen) ---
+        self.dino_model = DINOv3ViTModel.from_pretrained(dino_model_name)
+        self.dino_model.eval()
+        self.dino_model.requires_grad_(False)
+        
+        self.dino_transform = transforms.Compose([
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ])
+        
+        self.patch_size = self.dino_model.config.patch_size
+        self.patch_number = image_size // self.patch_size
+        self.embed_dim = self.dino_model.config.hidden_size  # e.g. 1024
+        
+        # --- Flux VAE encoder (frozen, lazy-loaded) ---
+        self.vae_model_name = vae_model_name
+        self._vae = None
+        self.vae_channels = 16  # Flux VAE outputs 16 channels
+        self.vae_downsample = 8  # Flux VAE downsamples by 8x
+        
+        # --- Projection grid (shared) ---
+        self.proj_grid = ProjGrid(
+            grid_resolution=grid_resolution,
+            image_resolution=image_size,
+        )
+        
+        # Expose dimensions for denoiser block construction
+        self.dino_proj_channels = self.embed_dim      # e.g. 1024
+        self.vae_proj_channels = self.vae_channels     # 16
+        # proj_channels is kept for backward compat with _proj_channels in mixin
+        self.proj_channels = self.embed_dim
+        
+    def _load_vae(self):
+        """Lazy-load Flux VAE encoder."""
+        if self._vae is not None:
+            return
+        from diffusers import AutoencoderKL
+        device = next(self.dino_model.parameters()).device
+        vae = AutoencoderKL.from_pretrained(
+            self.vae_model_name,
+            subfolder="vae",
+            torch_dtype=torch.float32,
+        )
+        vae.eval()
+        vae.requires_grad_(False)
+        vae.to(device)
+        self._vae = vae
+    
+    def to(self, device):
+        super().to(device)
+        self.dino_model.to(device)
+        self.proj_grid.to(device)
+        if self._vae is not None:
+            self._vae.to(device)
+        return self
+    
+    def cuda(self):
+        super().cuda()
+        self.dino_model.cuda()
+        self.proj_grid.cuda()
+        if self._vae is not None:
+            self._vae.cuda()
+        return self
+    
+    def cpu(self):
+        super().cpu()
+        self.dino_model.cpu()
+        self.proj_grid.cpu()
+        if self._vae is not None:
+            self._vae.cpu()
+        return self
+    
+    def _extract_dino_features(self, image: torch.Tensor) -> torch.Tensor:
+        """Extract DINOv3 features from normalized image."""
+        image = image.to(self.dino_model.embeddings.patch_embeddings.weight.dtype)
+        hidden_states = self.dino_model.embeddings(image, bool_masked_pos=None)
+        position_embeddings = self.dino_model.rope_embeddings(image)
+        for layer_module in self.dino_model.layer:
+            hidden_states = layer_module(
+                hidden_states,
+                position_embeddings=position_embeddings,
+            )
+        return F.layer_norm(hidden_states, hidden_states.shape[-1:])
+    
+    @torch.no_grad()
+    def _extract_vae_latent(self, image: torch.Tensor) -> torch.Tensor:
+        """Extract Flux VAE latent from unnormalized image [0,1]."""
+        self._load_vae()
+        image_normalized = image * 2.0 - 1.0
+        image_normalized = image_normalized.to(self._vae.dtype)
+        posterior = self._vae.encode(image_normalized)
+        latent = posterior.latent_dist.mode()
+        latent = latent * self._vae.config.scaling_factor
+        return latent.float()  # [B, 16, H/8, W/8]
+    
+    def forward(
+        self,
+        image: Union[torch.Tensor, List[Image.Image]],
+        camera_angle_x: Optional[torch.Tensor] = None,
+        distance: Optional[torch.Tensor] = None,
+        mesh_scale: Optional[torch.Tensor] = None,
+        transform_matrix: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Extract gated features from the image.
+        
+        Returns:
+            Tuple of (global_features, proj_semantic, proj_color):
+            - global_features: [B, num_global_tokens, embed_dim] (DINOv3 CLS + registers)
+            - proj_semantic: [B, grid_res³, embed_dim] (DINOv3 projected features)
+            - proj_color: [B, grid_res³, vae_channels] (VAE projected features)
+        """
+        # Handle input types
+        if isinstance(image, torch.Tensor):
+            assert image.ndim == 4
+        elif isinstance(image, list):
+            assert all(isinstance(i, Image.Image) for i in image)
+            image = [i.resize((self.image_size, self.image_size), Image.LANCZOS) for i in image]
+            image = [np.array(i.convert('RGB')).astype(np.float32) / 255 for i in image]
+            image = [torch.from_numpy(i).permute(2, 0, 1).float() for i in image]
+            image = torch.stack(image).cuda()
+        else:
+            raise ValueError(f"Unsupported type of image: {type(image)}")
+        
+        B = image.shape[0]
+        image_raw = image.clone()  # Keep unnormalized copy for VAE
+        
+        if camera_angle_x is None or distance is None or mesh_scale is None:
+            raise ValueError("camera_angle_x, distance, and mesh_scale must be provided")
+        
+        with torch.no_grad():
+            # --- DINOv3 branch ---
+            dino_input = self.dino_transform(image)
+            z = self._extract_dino_features(dino_input)
+            
+            z_clstoken = z[:, 0:1]
+            num_reg = getattr(self.dino_model.config, 'num_register_tokens', 4)
+            z_regtokens = z[:, 1:1+num_reg]
+            z_patchtokens = z[:, 1+num_reg:]
+            
+            z_patchtokens_spatial = z_patchtokens.reshape(
+                B, self.patch_number, self.patch_number, -1
+            )  # [B, h, w, D]
+            
+            proj_semantic = self.proj_grid(
+                z_patchtokens_spatial,
+                camera_angle_x, distance, mesh_scale, transform_matrix,
+            )  # [B, R³, embed_dim]
+            
+            z_global = torch.cat([z_clstoken, z_regtokens], dim=1)  # [B, 1+num_reg, D]
+            
+            # --- VAE branch ---
+            vae_latent = self._extract_vae_latent(image_raw)  # [B, 16, H/8, W/8]
+            
+            proj_color = self.proj_grid(
+                vae_latent,
+                camera_angle_x, distance, mesh_scale, transform_matrix,
+                BHWC=False,  # VAE latent is [B, C, H, W]
+            )  # [B, R³, 16]
+        
+        return z_global, proj_semantic, proj_color
+
+
+# =============================================================================
+# Image Conditioned Mixin with Projection Support
+# =============================================================================
+
+class ImageConditionedProjMixin:
+    """
+    Mixin for image-conditioned models with view-aligned projection.
+    
+    This mixin adds support for extracting view-aligned features from images
+    using camera parameters.
+    
+    Args:
+        image_cond_model: Configuration for the image conditioning model.
+    """
+    def __init__(self, *args, image_cond_model: dict, **kwargs):
+        # Store config before super().__init__ which calls init_models_and_more
+        self.image_cond_model_config = image_cond_model
+        self.image_cond_model = None  # Will be initialized in init_models_and_more
+        self.image_attn_mode = image_cond_model.get('image_attn_mode', 
+                                image_cond_model.get('args', {}).get('image_attn_mode', 'cross'))
+        super().__init__(*args, **kwargs)
+        
+    def _init_image_cond_model(self):
+        """Initialize the image conditioning model."""
+        with dist_utils.local_master_first():
+            model_name = self.image_cond_model_config['name']
+            model_args = self.image_cond_model_config.get('args', {})
+            
+            if model_name == 'DinoV3ProjFeatureExtractor':
+                self.image_cond_model = DinoV3ProjFeatureExtractor(**model_args)
+            elif model_name == 'DinoV3VaeProjFeatureExtractor':
+                self.image_cond_model = DinoV3VaeProjFeatureExtractor(**model_args)
+            else:
+                # Fallback to standard extractors
+                from . import image_conditioned
+                self.image_cond_model = getattr(image_conditioned, model_name)(**model_args)
+            
+            self.image_cond_model.cuda()
+            
+            # Expose proj_channels for denoiser to know the correct proj_in_channels
+            if hasattr(self.image_cond_model, 'proj_channels'):
+                self._proj_channels = self.image_cond_model.proj_channels
+            else:
+                self._proj_channels = getattr(self.image_cond_model, 'embed_dim', None)
+            # Expose vae_proj_channels for gated_proj mode
+            self._vae_proj_channels = getattr(self.image_cond_model, 'vae_proj_channels', None)
+    
+    def init_models_and_more(self, **kwargs):
+        """
+        Override to handle image_cond_model initialization.
+        
+        Since proj_linear has been moved to per-block ProjectAttention in the denoiser,
+        image_cond_model no longer has any trainable parameters (DINOv3 backbone is frozen,
+        ProjGrid only has register_buffers). Therefore we do NOT add it to self.models
+        (which would trigger DDP wrapping and fail). We just initialize it and keep it
+        as a standalone module for inference.
+        """
+        # Initialize image_cond_model first
+        if self.image_cond_model is None:
+            self._init_image_cond_model()
+        
+        # Keep a reference to the unwrapped module for attribute access
+        self._image_cond_module = self.image_cond_model  # for .grid_resolution etc.
+        
+        # Log that image_cond has no trainable params
+        proj_params = [p for p in self.image_cond_model.parameters() if p.requires_grad]
+        if self.is_master:
+            if proj_params:
+                print(f'\nWARNING: image_cond_model has {len(proj_params)} trainable params, '
+                      f'but is NOT registered in self.models. These will NOT be trained!')
+            else:
+                print(f'\nimage_cond_model has no trainable parameters, skipping DDP/optimizer registration.')
+        
+        # Call base class to set up DDP, optimizer, EMA, etc. (without image_cond)
+        super().init_models_and_more(**kwargs)
+
+    # ------------------------------------------------------------------
+    # Checkpoint save/load overrides: skip DINOv3 backbone weights
+    # ------------------------------------------------------------------
+
+    # Keys in image_cond state_dict that belong to the frozen DINOv3 backbone.
+    # Everything under "model." is DINOv3; we only keep proj_grid.*
+    _IMAGE_COND_BACKBONE_PREFIX = 'model.'
+
+    def _filter_image_cond_state_dict(self, state_dict: dict) -> dict:
+        """Keep only non-backbone keys (proj_grid, etc.) from image_cond state_dict."""
+        return {k: v for k, v in state_dict.items()
+                if not k.startswith(self._IMAGE_COND_BACKBONE_PREFIX)}
+
+    def _fill_denoiser_proj_linear_from_image_cond(
+        self,
+        denoiser_ckpt: dict,
+        denoiser_state_dict: dict,
+        image_cond_ckpt_path: Optional[str] = None,
+    ) -> dict:
+        """
+        Fill missing per-block proj_linear weights in denoiser checkpoint
+        from the old-style image_cond proj_linear (broadcast to all blocks).
+        
+        Also handles shape mismatch when NAF is enabled: old proj_linear has shape
+        [model_ch, embed_dim] but new model expects [model_ch, embed_dim*2].
+        In this case, the old weights are placed in the lr half and the hr half is zero-padded.
+        
+        Compatibility strategy:
+        1. If denoiser_ckpt already contains per-block proj_linear keys with correct shape -> do nothing.
+        2. If shape mismatch (embed_dim vs embed_dim*2) -> zero-pad the weight.
+        3. If keys missing, try to load proj_linear from image_cond checkpoint -> broadcast (with optional pad).
+        
+        Args:
+            denoiser_ckpt: The loaded denoiser state dict
+            denoiser_state_dict: The model's current state dict (to find expected keys)
+            image_cond_ckpt_path: Path to image_cond checkpoint file (optional)
+            
+        Returns:
+            Updated denoiser_ckpt with proj_linear keys filled if needed
+        """
+        if self.image_attn_mode != 'proj':
+            return denoiser_ckpt
+        
+        # Find all per-block proj_linear keys expected by the model
+        proj_linear_keys = [k for k in denoiser_state_dict.keys()
+                            if '.cross_attn.proj_linear.' in k]
+        if not proj_linear_keys:
+            return denoiser_ckpt
+        
+        # --- Phase 1: Handle shape mismatch for existing keys (NAF upgrade) ---
+        for k in proj_linear_keys:
+            if k in denoiser_ckpt:
+                expected_shape = denoiser_state_dict[k].shape
+                actual_shape = denoiser_ckpt[k].shape
+                if expected_shape != actual_shape:
+                    if k.endswith('.weight') and len(expected_shape) == 2:
+                        # Weight shape: [out_features, in_features]
+                        # Old: [model_ch, embed_dim], New: [model_ch, embed_dim*2]
+                        out_f, new_in_f = expected_shape
+                        _, old_in_f = actual_shape
+                        if new_in_f > old_in_f and out_f == actual_shape[0]:
+                            if self.is_master:
+                                print(f'\n  [NAF Compat] Padding proj_linear weight {k}: '
+                                      f'{actual_shape} -> {expected_shape} (zero-pad hr half)')
+                            new_w = torch.zeros(expected_shape, dtype=denoiser_ckpt[k].dtype,
+                                                device=denoiser_ckpt[k].device)
+                            new_w[:, :old_in_f] = denoiser_ckpt[k]
+                            denoiser_ckpt[k] = new_w
+                        else:
+                            if self.is_master:
+                                print(f'\n  Warning: proj_linear {k} shape mismatch '
+                                      f'{actual_shape} vs {expected_shape}, using model init')
+                            denoiser_ckpt[k] = denoiser_state_dict[k]
+                    # bias shape should match (out_features only), no padding needed
+        
+        # --- Phase 2: Handle completely missing keys ---
+        missing_proj_keys = [k for k in proj_linear_keys if k not in denoiser_ckpt]
+        if not missing_proj_keys:
+            return denoiser_ckpt
+        
+        if self.is_master:
+            print(f'\n  [Compat] Denoiser ckpt missing {len(missing_proj_keys)} per-block proj_linear keys.')
+            print(f'           Attempting to load from image_cond proj_linear: {image_cond_ckpt_path}')
+        
+        # Try to find proj_linear weights from image_cond checkpoint
+        old_proj_linear_w = None
+        old_proj_linear_b = None
+        
+        if image_cond_ckpt_path is not None:
+            import os as _os
+            if _os.path.exists(image_cond_ckpt_path):
+                try:
+                    ic_ckpt = torch.load(image_cond_ckpt_path, map_location=self.device, weights_only=True)
+                    old_proj_linear_w = ic_ckpt.get('proj_linear.weight')
+                    old_proj_linear_b = ic_ckpt.get('proj_linear.bias')
+                except Exception as e:
+                    if self.is_master:
+                        print(f'           Warning: Failed to load image_cond ckpt: {e}')
+        
+        if old_proj_linear_w is None:
+            raise RuntimeError(
+                f'Denoiser checkpoint is missing per-block proj_linear keys '
+                f'(e.g. {missing_proj_keys[0]}), and no image_cond proj_linear '
+                f'was found to broadcast from. Cannot proceed.'
+            )
+        
+        if self.is_master:
+            print(f'           Found image_cond proj_linear: weight {old_proj_linear_w.shape}, bias {old_proj_linear_b.shape}')
+            print(f'           Broadcasting to {len(missing_proj_keys)} per-block keys...')
+        
+        for k in missing_proj_keys:
+            if k.endswith('.weight'):
+                expected_shape = denoiser_state_dict[k].shape
+                if expected_shape != old_proj_linear_w.shape:
+                    # Pad for NAF: [model_ch, embed_dim] -> [model_ch, embed_dim*2]
+                    out_f, new_in_f = expected_shape
+                    _, old_in_f = old_proj_linear_w.shape
+                    if new_in_f > old_in_f and out_f == old_proj_linear_w.shape[0]:
+                        new_w = torch.zeros(expected_shape, dtype=old_proj_linear_w.dtype)
+                        new_w[:, :old_in_f] = old_proj_linear_w
+                        denoiser_ckpt[k] = new_w
+                    else:
+                        denoiser_ckpt[k] = denoiser_state_dict[k]
+                else:
+                    denoiser_ckpt[k] = old_proj_linear_w.clone()
+            elif k.endswith('.bias'):
+                denoiser_ckpt[k] = old_proj_linear_b.clone()
+        
+        return denoiser_ckpt
+
+    def _master_params_to_state_dicts(self, master_params):
+        """Override to skip image_cond checkpoint entirely.
+        
+        image_cond model no longer has trainable parameters:
+        - proj_linear has been moved to per-block ProjectAttention in the denoiser
+        - DINOv3 backbone is frozen and loaded from pretrained weights
+        - ProjGrid only contains fixed register_buffers (grid_points, front_view_transform_matrix)
+        So there is nothing worth saving for image_cond.
+        """
+        state_dicts = super()._master_params_to_state_dicts(master_params)
+        state_dicts.pop('image_cond', None)
+        return state_dicts
+
+    def load(self, load_dir, step=0):
+        """
+        Override to handle:
+        1. Old checkpoints that don't have image_cond_step*.pt
+        2. Partial image_cond checkpoints (only proj_linear + proj_grid, no DINOv3 backbone)
+        """
+        import os as _os
+
+        if self.is_master:
+            print(f'\nLoading checkpoint from step {step}...', end='')
+
+        model_ckpts = {}
+        for name, model in self.models.items():
+            ckpt_path = _os.path.join(load_dir, 'ckpts', f'{name}_step{step:07d}.pt')
+
+            if name == 'image_cond':
+                # --- handle missing or partial image_cond checkpoint ---
+                if not _os.path.exists(ckpt_path):
+                    if self.is_master:
+                        print(f'\n  image_cond checkpoint not found at {ckpt_path}, using freshly initialised weights.')
+                    model_ckpts[name] = model.state_dict()
+                    continue
+
+                try:
+                    model_ckpt = torch.load(
+                        read_file_dist(ckpt_path),
+                        map_location=self.device, weights_only=True)
+                except Exception as e:
+                    if self.is_master:
+                        print(f'\n  Failed to load image_cond checkpoint: {e}. Using freshly initialised weights.')
+                    model_ckpts[name] = model.state_dict()
+                    continue
+
+                # Partial ckpt (no backbone) → load with strict=False
+                missing, unexpected = model.load_state_dict(model_ckpt, strict=False)
+                # All missing keys should be the frozen DINOv3 backbone; verify
+                non_backbone_missing = [k for k in missing
+                                        if not k.startswith(self._IMAGE_COND_BACKBONE_PREFIX)]
+                if non_backbone_missing and self.is_master:
+                    print(f'\n  Warning: unexpected missing keys in image_cond ckpt: {non_backbone_missing}')
+                if unexpected and self.is_master:
+                    print(f'\n  Warning: unexpected keys in image_cond ckpt: {unexpected}')
+
+                # Build a full state_dict for master_params sync
+                full_sd = model.state_dict()
+                full_sd.update(model_ckpt)
+                model_ckpts[name] = full_sd
+            else:
+                model_ckpt = torch.load(
+                    read_file_dist(ckpt_path),
+                    map_location=self.device, weights_only=True)
+                # For denoiser: handle old ckpts missing per-block proj_linear
+                if name == 'denoiser':
+                    ic_ckpt_path = _os.path.join(load_dir, 'ckpts', f'image_cond_step{step:07d}.pt')
+                    model_ckpt = self._fill_denoiser_proj_linear_from_image_cond(
+                        model_ckpt, model.state_dict(), ic_ckpt_path)
+                model_ckpts[name] = model_ckpt
+                model.load_state_dict(model_ckpt)
+
+        self._state_dicts_to_master_params(self.master_params, model_ckpts)
+        del model_ckpts
+
+        if self.is_master:
+            for i, ema_rate in enumerate(self.ema_rate):
+                ema_ckpts = {}
+                for name, model in self.models.items():
+                    ema_path = _os.path.join(
+                        load_dir, 'ckpts',
+                        f'{name}_ema{ema_rate}_step{step:07d}.pt')
+                    if name == 'image_cond':
+                        if not _os.path.exists(ema_path):
+                            ema_ckpts[name] = model.state_dict()
+                            continue
+                        try:
+                            ema_ckpt = torch.load(ema_path, map_location=self.device, weights_only=True)
+                        except Exception:
+                            ema_ckpts[name] = model.state_dict()
+                            continue
+                        full_sd = model.state_dict()
+                        full_sd.update(ema_ckpt)
+                        ema_ckpts[name] = full_sd
+                    else:
+                        ema_ckpt = torch.load(ema_path, map_location=self.device, weights_only=True)
+                        if name == 'denoiser':
+                            ic_ema_path = _os.path.join(
+                                load_dir, 'ckpts',
+                                f'image_cond_ema{ema_rate}_step{step:07d}.pt')
+                            ema_ckpt = self._fill_denoiser_proj_linear_from_image_cond(
+                                ema_ckpt, model.state_dict(), ic_ema_path)
+                        ema_ckpts[name] = ema_ckpt
+                self._state_dicts_to_master_params(self.ema_params[i], ema_ckpts)
+                del ema_ckpts
+
+        misc_ckpt = torch.load(
+            read_file_dist(_os.path.join(load_dir, 'ckpts', f'misc_step{step:07d}.pt')),
+            map_location=torch.device('cpu'), weights_only=False)
+        # Optimizer state may mismatch when loading old checkpoints that were
+        # saved before image_cond was added to self.models, or when the number
+        # of trainable parameters changed (e.g. backbone freeze, NAF upgrade).
+        # In that case we skip restoring optimizer state and let it re-initialise.
+        try:
+            self.optimizer.load_state_dict(misc_ckpt['optimizer'])
+            # Verify optimizer state shapes match parameters.
+            # load_state_dict may succeed even when shapes mismatch (keys are
+            # integer indices), causing a crash later in optimizer.step().
+            _shape_ok = True
+            for group in self.optimizer.param_groups:
+                for p in group['params']:
+                    state = self.optimizer.state.get(p)
+                    if state is not None:
+                        for sv in state.values():
+                            if isinstance(sv, torch.Tensor) and sv.shape != () and sv.shape != p.shape:
+                                _shape_ok = False
+                                break
+                    if not _shape_ok:
+                        break
+                if not _shape_ok:
+                    break
+            if not _shape_ok:
+                if self.is_master:
+                    print(f'\n  Warning: optimizer state shape mismatch (likely NAF upgrade). '
+                          f'Optimizer will start fresh.')
+                self.optimizer.state.clear()
+        except (ValueError, RuntimeError) as e:
+            if self.is_master:
+                print(f'\n  Warning: could not load optimizer state ({e}). '
+                      f'Optimizer will start fresh.')
+        self.step = misc_ckpt['step']
+        self.data_sampler.load_state_dict(misc_ckpt['data_sampler'])
+        if self.mix_precision_mode == 'amp' and self.mix_precision_dtype == torch.float16:
+            self.scaler.load_state_dict(misc_ckpt['scaler'])
+        elif self.mix_precision_mode == 'inflat_all' and self.mix_precision_dtype == torch.float16:
+            self.log_scale = misc_ckpt['log_scale']
+        if self.lr_scheduler_config is not None:
+            self.lr_scheduler.load_state_dict(misc_ckpt['lr_scheduler'])
+        if self.elastic_controller_config is not None:
+            self.elastic_controller.load_state_dict(misc_ckpt['elastic_controller'])
+        if self.grad_clip is not None and not isinstance(self.grad_clip, float):
+            self.grad_clip.load_state_dict(misc_ckpt['grad_clip'])
+        del misc_ckpt
+
+        if self.world_size > 1:
+            dist.barrier()
+        if self.is_master:
+            print(' Done.')
+
+        if self.world_size > 1:
+            self.check_ddp()
+
+    def finetune_from(self, finetune_ckpt):
+        """
+        Override to tolerate DINOv3 backbone keys missing from image_cond checkpoint.
+        For image_cond, the checkpoint only stores proj_linear + proj_grid (no backbone),
+        so we treat all backbone keys as allowed-missing.
+        """
+        ALLOWED_MISSING_KEYS = {'rope_phases'}
+
+        if self.is_master:
+            print('\nFinetuning from:')
+            for name, path in finetune_ckpt.items():
+                print(f'  - {name}: {path}')
+
+        model_ckpts = {}
+        for name, model in self.models.items():
+            model_state_dict = model.state_dict()
+            if name in finetune_ckpt:
+                model_ckpt = torch.load(
+                    read_file_dist(finetune_ckpt[name]),
+                    map_location=self.device, weights_only=True)
+
+                model_ckpt = self._remap_checkpoint_keys(model_ckpt, model_state_dict)
+
+                for k, v in model_ckpt.items():
+                    if k not in model_state_dict:
+                        if self.is_master:
+                            print(f'Warning: {k} not found in model_state_dict, skipped.')
+                        model_ckpt[k] = None
+                    elif model_ckpt[k].shape != model_state_dict[k].shape:
+                        # For proj_linear weights, try zero-pad instead of skipping
+                        # This handles NAF upgrade: [model_ch, embed_dim] -> [model_ch, embed_dim*2]
+                        if '.cross_attn.proj_linear.weight' in k and len(model_ckpt[k].shape) == 2:
+                            old_shape = model_ckpt[k].shape
+                            new_shape = model_state_dict[k].shape
+                            if new_shape[0] == old_shape[0] and new_shape[1] > old_shape[1]:
+                                if self.is_master:
+                                    print(f'Info: Zero-padding proj_linear weight {k}: {old_shape} -> {new_shape}')
+                                new_w = torch.zeros(new_shape, dtype=model_ckpt[k].dtype)
+                                new_w[:, :old_shape[1]] = model_ckpt[k]
+                                model_ckpt[k] = new_w
+                            else:
+                                if self.is_master:
+                                    print(f'Warning: {k} shape mismatch, {old_shape} vs {new_shape}, skipped.')
+                                model_ckpt[k] = model_state_dict[k]
+                        else:
+                            if self.is_master:
+                                print(f'Warning: {k} shape mismatch, {model_ckpt[k].shape} vs {model_state_dict[k].shape}, skipped.')
+                            model_ckpt[k] = model_state_dict[k]
+                model_ckpt = {k: v for k, v in model_ckpt.items() if v is not None}
+
+                missing_keys = set(model_state_dict.keys()) - set(model_ckpt.keys())
+
+                # For denoiser: fill per-block proj_linear from image_cond if missing
+                if name == 'denoiser':
+                    ic_path = finetune_ckpt.get('image_cond')
+                    proj_linear_missing = {k for k in missing_keys if '.cross_attn.proj_linear.' in k}
+                    if proj_linear_missing:
+                        model_ckpt = self._fill_denoiser_proj_linear_from_image_cond(
+                            model_ckpt, model_state_dict, ic_path)
+                        # Recalculate missing_keys after filling
+                        missing_keys = set(model_state_dict.keys()) - set(model_ckpt.keys())
+
+                # For image_cond, DINOv3 backbone keys are expected to be missing
+                allowed = set(ALLOWED_MISSING_KEYS)
+                if name == 'image_cond':
+                    backbone_missing = {k for k in missing_keys
+                                        if k.startswith(self._IMAGE_COND_BACKBONE_PREFIX)}
+                    allowed |= backbone_missing
+                    if backbone_missing and self.is_master:
+                        print(f'Info: image_cond: {len(backbone_missing)} DINOv3 backbone keys '
+                              f'not in ckpt (expected, using pretrained weights)')
+                    # Old ckpts may have proj_linear.* which has moved to denoiser blocks
+                    proj_linear_missing = {k for k in missing_keys if k.startswith('proj_linear.')}
+                    allowed |= proj_linear_missing
+
+                unexpected_missing = missing_keys - allowed
+                if unexpected_missing and self.is_master:
+                    print(f'Error: Missing keys in checkpoint: {unexpected_missing}')
+                    raise RuntimeError(f'Missing keys in checkpoint: {unexpected_missing}')
+                if missing_keys & ALLOWED_MISSING_KEYS and self.is_master:
+                    print(f'Info: Using model initialized values for: {missing_keys & ALLOWED_MISSING_KEYS}')
+
+                for k in missing_keys:
+                    model_ckpt[k] = model_state_dict[k]
+
+                model_ckpts[name] = model_ckpt
+                model.load_state_dict(model_ckpt)
+            else:
+                if self.is_master:
+                    print(f'Warning: {name} not found in finetune_ckpt, skipped.')
+                model_ckpts[name] = model_state_dict
+
+        self._state_dicts_to_master_params(self.master_params, model_ckpts)
+        if self.is_master:
+            for i, ema_rate in enumerate(self.ema_rate):
+                self._state_dicts_to_master_params(self.ema_params[i], model_ckpts)
+        del model_ckpts
+
+        if self.world_size > 1:
+            dist.barrier()
+        if self.is_master:
+            print('Done.')
+
+        if self.world_size > 1:
+            self.check_ddp()
+
+    def encode_image_proj(
+        self,
+        image: torch.Tensor,
+        camera_angle_x: Optional[torch.Tensor] = None,
+        distance: Optional[torch.Tensor] = None,
+        mesh_scale: Optional[torch.Tensor] = None,
+        transform_matrix: Optional[torch.Tensor] = None,
+        coords: Optional[torch.Tensor] = None,
+    ) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor]]:
+        """
+        Encode the image with view-aligned projection.
+        
+        Supports both 'proj' mode (DINOv3 only, 2 outputs) and 
+        'gated_proj' mode (DINOv3 + VAE, 3 outputs).
+        """
+        if self.image_cond_model is None:
+            self._init_image_cond_model()
+        
+        outputs = self.image_cond_model(
+            image,
+            camera_angle_x=camera_angle_x,
+            distance=distance,
+            mesh_scale=mesh_scale,
+            transform_matrix=transform_matrix,
+        )
+        
+        is_gated = self.image_attn_mode == 'gated_proj'
+        
+        if is_gated:
+            cond_global, cond_proj_semantic, cond_proj_color = outputs
+        else:
+            cond_global, cond_proj = outputs
+        
+        # If coords provided, extract features at sparse positions
+        if coords is not None:
+            B = cond_global.shape[0]
+            module = getattr(self, '_image_cond_module', self.image_cond_model)
+            grid_res = module.grid_resolution
+            batch_indices = coords[:, 0].long()
+            x_coords = coords[:, 1].long()
+            y_coords = coords[:, 2].long()
+            z_coords = coords[:, 3].long()
+            
+            if is_gated:
+                cond_proj_semantic = cond_proj_semantic.reshape(B, grid_res, grid_res, grid_res, -1)
+                cond_proj_semantic = cond_proj_semantic[batch_indices, x_coords, y_coords, z_coords]
+                cond_proj_color = cond_proj_color.reshape(B, grid_res, grid_res, grid_res, -1)
+                cond_proj_color = cond_proj_color[batch_indices, x_coords, y_coords, z_coords]
+            else:
+                cond_proj = cond_proj.reshape(B, grid_res, grid_res, grid_res, -1)
+                cond_proj = cond_proj[batch_indices, x_coords, y_coords, z_coords]
+        
+        if is_gated:
+            cond = {
+                'global': cond_global,
+                'proj_semantic': cond_proj_semantic,
+                'proj_color': cond_proj_color,
+            }
+            uncond = {
+                'global': torch.zeros_like(cond_global),
+                'proj_semantic': torch.zeros_like(cond_proj_semantic),
+                'proj_color': torch.zeros_like(cond_proj_color),
+            }
+        else:
+            cond = {'global': cond_global, 'proj': cond_proj}
+            uncond = {'global': torch.zeros_like(cond_global), 'proj': torch.zeros_like(cond_proj)}
+        
+        return cond, uncond
+    
+    @torch.no_grad()
+    def encode_image(self, image: Union[torch.Tensor, List[Image.Image]]) -> torch.Tensor:
+        """
+        Encode the image (standard mode without projection).
+        """
+        if self.image_cond_model is None:
+            self._init_image_cond_model()
+        
+        if self.image_attn_mode == 'proj':
+            # For proj mode, return dict
+            global_feat, proj_feat = self.image_cond_model(image)
+            return {'global': global_feat, 'proj': proj_feat}
+        else:
+            # Standard mode
+            features = self.image_cond_model(image)
+            return features
+        
+    def _extract_camera_info(self, kwargs):
+        """
+        Extract camera info from kwargs.
+        
+        Supports two formats:
+        1. 'camera_info' dict: {'camera_angle_x': ..., 'distance': ..., 'mesh_scale': ..., 'transform_matrix': ..., 'coords': ...}
+        2. Flat fields: 'camera_angle_x', 'camera_distance', 'mesh_scale', 'transform_matrix', 'coords' in kwargs
+        
+        Returns:
+            camera_info dict or None if not available
+        """
+        if 'camera_info' in kwargs:
+            return kwargs.pop('camera_info')
+        
+        # Try to extract from flat fields (as returned by ViewImageConditionedMixin)
+        camera_angle_x = kwargs.pop('camera_angle_x', None)
+        camera_distance = kwargs.pop('camera_distance', None)
+        mesh_scale = kwargs.pop('mesh_scale', None)
+        transform_matrix = kwargs.pop('transform_matrix', None)
+        coords = kwargs.pop('coords', None)
+        
+        if camera_angle_x is not None and camera_distance is not None and mesh_scale is not None:
+            return {
+                'camera_angle_x': camera_angle_x,
+                'distance': camera_distance,
+                'mesh_scale': mesh_scale,
+                'transform_matrix': transform_matrix,
+                'coords': coords,
+            }
+        
+        return None
+        
+    def get_cond(self, cond, **kwargs):
+        """Get the conditioning data."""
+        kwargs.pop('view_idx', None)
+        
+        if self.image_attn_mode in ('proj', 'gated_proj'):
+            # Handle projection mode (both standard proj and gated_proj)
+            camera_info = self._extract_camera_info(kwargs)
+            if camera_info is not None:
+                coords = camera_info.get('coords')
+                cond, neg_cond = self.encode_image_proj(
+                    cond,
+                    camera_angle_x=camera_info.get('camera_angle_x'),
+                    distance=camera_info.get('distance'),
+                    mesh_scale=camera_info.get('mesh_scale'),
+                    transform_matrix=camera_info.get('transform_matrix'),
+                    coords=coords,
+                )
+                
+                # For sparse mode (coords provided), handle CFG dropout ourselves
+                if coords is not None and hasattr(self, 'p_uncond') and self.p_uncond > 0:
+                    import numpy as np
+                    B = cond['global'].shape[0]
+                    mask = np.random.rand(B) < self.p_uncond
+                    
+                    global_tensor = cond['global']
+                    global_mask_shape = [B] + [1] * (global_tensor.ndim - 1)
+                    global_mask = torch.tensor(mask, device=global_tensor.device).reshape(global_mask_shape)
+                    cond['global'] = torch.where(global_mask, neg_cond['global'], cond['global'])
+                    
+                    batch_indices = coords[:, 0].long()
+                    # Handle all sparse proj keys (proj, or proj_semantic + proj_color)
+                    for key in list(cond.keys()):
+                        if key.startswith('proj'):
+                            device = cond[key].device
+                            sparse_mask = torch.tensor(mask, device=device)[batch_indices].reshape(-1, 1)
+                            cond[key] = torch.where(sparse_mask, neg_cond[key], cond[key])
+                    
+                    return cond
+                else:
+                    kwargs['neg_cond'] = neg_cond
+            else:
+                cond = self.encode_image(cond)
+                if isinstance(cond, dict) and 'global' in cond:
+                    kwargs['neg_cond'] = {k: torch.zeros_like(v) for k, v in cond.items()}
+                else:
+                    kwargs['neg_cond'] = torch.zeros_like(cond)
+        else:
+            cond = self.encode_image(cond)
+            kwargs['neg_cond'] = torch.zeros_like(cond)
+            
+        cond = super().get_cond(cond, **kwargs)
+        return cond
+    
+    def get_inference_cond(self, cond, **kwargs):
+        """Get the conditioning data for inference."""
+        kwargs.pop('view_idx', None)
+        
+        if self.image_attn_mode in ('proj', 'gated_proj'):
+            camera_info = self._extract_camera_info(kwargs)
+            if camera_info is not None:
+                cond, neg_cond = self.encode_image_proj(
+                    cond,
+                    camera_angle_x=camera_info.get('camera_angle_x'),
+                    distance=camera_info.get('distance'),
+                    mesh_scale=camera_info.get('mesh_scale'),
+                    transform_matrix=camera_info.get('transform_matrix'),
+                    coords=camera_info.get('coords'),
+                )
+                kwargs['neg_cond'] = neg_cond
+            else:
+                cond = self.encode_image(cond)
+                if isinstance(cond, dict) and 'global' in cond:
+                    kwargs['neg_cond'] = {k: torch.zeros_like(v) for k, v in cond.items()}
+                else:
+                    kwargs['neg_cond'] = torch.zeros_like(cond)
+        else:
+            cond = self.encode_image(cond)
+            kwargs['neg_cond'] = torch.zeros_like(cond)
+            
+        cond = super().get_inference_cond(cond, **kwargs)
+        return cond
+
+    def vis_cond(self, cond, **kwargs):
+        """Visualize the conditioning data."""
+        return {'image': {'value': cond, 'type': 'image'}}
+
+    @torch.no_grad()
+    def visualize_projection_test(
+        self,
+        cond: torch.Tensor,
+        save_dir: str,
+        prefix: str = "proj_vis",
+        **kwargs
+    ) -> Optional[List[Image.Image]]:
+        """
+        Visualize projection points on the condition images.
+        
+        This should be called once before training starts to verify the projection is correct.
+        
+        Args:
+            cond: Condition image tensor [B, C, H, W], in [0, 1] range
+            save_dir: Directory to save visualizations
+            prefix: Prefix for saved files
+            **kwargs: Should contain camera_angle_x, camera_distance, mesh_scale, transform_matrix
+            
+        Returns:
+            List of PIL Images with projected points overlaid, or None if not in proj mode
+        """
+        if self.image_attn_mode != 'proj':
+            return None
+            
+        if self.image_cond_model is None:
+            self._init_image_cond_model()
+        
+        # Use _image_cond_module for attribute access (image_cond_model may be DDP-wrapped)
+        module = getattr(self, '_image_cond_module', self.image_cond_model)
+        
+        # Check if the model has visualization capability
+        if not hasattr(module, 'visualize_projection'):
+            print("Warning: image_cond_model does not support visualize_projection")
+            return None
+        
+        # Extract camera info
+        camera_info = self._extract_camera_info(kwargs)
+        if camera_info is None:
+            print("Warning: No camera info available for projection visualization")
+            return None
+        
+        return module.visualize_projection(
+            image=cond,
+            camera_angle_x=camera_info.get('camera_angle_x'),
+            distance=camera_info.get('distance'),
+            mesh_scale=camera_info.get('mesh_scale'),
+            transform_matrix=camera_info.get('transform_matrix'),
+            save_dir=save_dir,
+            prefix=prefix,
+        )
diff --git a/pixal3d/trainers/flow_matching/mixins/text_conditioned.py b/pixal3d/trainers/flow_matching/mixins/text_conditioned.py
new file mode 100644
index 0000000..85f1dcf
--- /dev/null
+++ b/pixal3d/trainers/flow_matching/mixins/text_conditioned.py
@@ -0,0 +1,68 @@
+from typing import *
+import os
+os.environ['TOKENIZERS_PARALLELISM'] = 'true'
+import torch
+from transformers import AutoTokenizer, CLIPTextModel
+
+from ....utils import dist_utils
+
+
+class TextConditionedMixin:
+    """
+    Mixin for text-conditioned models.
+    
+    Args:
+        text_cond_model: The text conditioning model.
+    """
+    def __init__(self, *args, text_cond_model: str = 'openai/clip-vit-large-patch14', **kwargs):
+        super().__init__(*args, **kwargs)
+        self.text_cond_model_name = text_cond_model
+        self.text_cond_model = None     # the model is init lazily
+        
+    def _init_text_cond_model(self):
+        """
+        Initialize the text conditioning model.
+        """
+        # load model
+        with dist_utils.local_master_first():
+            model = CLIPTextModel.from_pretrained(self.text_cond_model_name)
+            tokenizer = AutoTokenizer.from_pretrained(self.text_cond_model_name)
+        model.eval()
+        model = model.cuda()
+        self.text_cond_model = {
+            'model': model,
+            'tokenizer': tokenizer,
+        }
+        self.text_cond_model['null_cond'] = self.encode_text([''])
+        
+    @torch.no_grad()
+    def encode_text(self, text: List[str]) -> torch.Tensor:
+        """
+        Encode the text.
+        """
+        assert isinstance(text, list) and isinstance(text[0], str), "TextConditionedMixin only supports list of strings as cond"
+        if self.text_cond_model is None:
+            self._init_text_cond_model()
+        encoding = self.text_cond_model['tokenizer'](text, max_length=77, padding='max_length', truncation=True, return_tensors='pt')
+        tokens = encoding['input_ids'].cuda()
+        embeddings = self.text_cond_model['model'](input_ids=tokens).last_hidden_state
+        
+        return embeddings
+        
+    def get_cond(self, cond, **kwargs):
+        """
+        Get the conditioning data.
+        """
+        cond = self.encode_text(cond)
+        kwargs['neg_cond'] = self.text_cond_model['null_cond'].repeat(cond.shape[0], 1, 1)
+        cond = super().get_cond(cond, **kwargs)
+        return cond
+    
+    def get_inference_cond(self, cond, **kwargs):
+        """
+        Get the conditioning data for inference.
+        """
+        cond = self.encode_text(cond)
+        kwargs['neg_cond'] = self.text_cond_model['null_cond'].repeat(cond.shape[0], 1, 1)
+        cond = super().get_inference_cond(cond, **kwargs)
+        return cond
diff --git a/pixal3d/trainers/flow_matching/sparse_flow_matching.py b/pixal3d/trainers/flow_matching/sparse_flow_matching.py
new file mode 100644
index 0000000..65f211f
--- /dev/null
+++ b/pixal3d/trainers/flow_matching/sparse_flow_matching.py
@@ -0,0 +1,615 @@
+from typing import *
+import os
+import copy
+import functools
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+import numpy as np
+from easydict import EasyDict as edict
+
+from ...modules import sparse as sp
+from ...utils.general_utils import dict_reduce
+from ...utils.data_utils import recursive_to_device, cycle, BalancedResumableSampler
+from .flow_matching import FlowMatchingTrainer
+from .mixins.classifier_free_guidance import ClassifierFreeGuidanceMixin
+from .mixins.text_conditioned import TextConditionedMixin
+from .mixins.image_conditioned import ImageConditionedMixin, MultiImageConditionedMixin
+from .mixins.image_conditioned_proj import ImageConditionedProjMixin
+
+
+class SparseFlowMatchingTrainer(FlowMatchingTrainer):
+    """
+    Trainer for sparse diffusion model with flow matching objective.
+    
+    Args:
+        models (dict[str, nn.Module]): Models to train.
+        dataset (torch.utils.data.Dataset): Dataset.
+        output_dir (str): Output directory.
+        load_dir (str): Load directory.
+        step (int): Step to load.
+        batch_size (int): Batch size.
+        batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
+        batch_split (int): Split batch with gradient accumulation.
+        max_steps (int): Max steps.
+        optimizer (dict): Optimizer config.
+        lr_scheduler (dict): Learning rate scheduler config.
+        elastic (dict): Elastic memory management config.
+        grad_clip (float or dict): Gradient clip config.
+        ema_rate (float or list): Exponential moving average rates.
+        fp16_mode (str): FP16 mode.
+            - None: No FP16.
+            - 'inflat_all': Hold a inflated fp32 master param for all params.
+            - 'amp': Automatic mixed precision.
+        fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
+        finetune_ckpt (dict): Finetune checkpoint.
+        log_param_stats (bool): Log parameter stats.
+        i_print (int): Print interval.
+        i_log (int): Log interval.
+        i_sample (int): Sample interval.
+        i_save (int): Save interval.
+        i_ddpcheck (int): DDP check interval.
+
+        t_schedule (dict): Time schedule for flow matching.
+        sigma_min (float): Minimum noise level.
+    """
+    
+    def prepare_dataloader(self, **kwargs):
+        """
+        Prepare dataloader.
+        """
+        self.data_sampler = BalancedResumableSampler(
+            self.dataset,
+            shuffle=True,
+            batch_size=self.batch_size_per_gpu,
+        )
+        if self.num_workers is None or self.num_workers == -1:
+            num_workers = max(1, int(np.ceil((os.cpu_count() - 16) / torch.cuda.device_count())))
+        else:
+            num_workers = self.num_workers
+        
+        self.dataloader = DataLoader(
+            self.dataset,
+            batch_size=self.batch_size_per_gpu,
+            num_workers=num_workers,
+            pin_memory=True,
+            drop_last=True,
+            persistent_workers=True,
+            collate_fn=functools.partial(self.dataset.collate_fn, split_size=self.batch_split),
+            sampler=self.data_sampler,
+        )
+        self.data_iterator = cycle(self.dataloader)
+
+    def training_losses(
+        self,
+        x_0: sp.SparseTensor,
+        cond=None,
+        **kwargs
+    ) -> Tuple[Dict, Dict]:
+        """
+        Compute training losses for a single timestep.
+
+        Args:
+            x_0: The [N x ... x C] sparse tensor of the inputs.
+            cond: The [N x ...] tensor of additional conditions.
+            kwargs: Additional arguments to pass to the backbone.
+
+        Returns:
+            a dict with the key "loss" containing a tensor of shape [N].
+            may also contain other keys for different terms.
+        """
+        noise = x_0.replace(torch.randn_like(x_0.feats))
+        t = self.sample_t(x_0.shape[0]).to(x_0.device).float()
+        x_t = self.diffuse(x_0, t, noise=noise)
+        cond = self.get_cond(cond, **kwargs)
+        
+        pred = self.training_models['denoiser'](x_t, t * 1000, cond, **kwargs)
+        assert pred.shape == noise.shape == x_0.shape
+        target = self.get_v(x_0, noise, t)
+        terms = edict()
+        terms["mse"] = F.mse_loss(pred.feats, target.feats)
+        terms["loss"] = terms["mse"]
+
+        # log loss with time bins
+        mse_per_instance = np.array([
+            F.mse_loss(pred.feats[x_0.layout[i]], target.feats[x_0.layout[i]]).item()
+            for i in range(x_0.shape[0])
+        ])
+        time_bin = np.digitize(t.cpu().numpy(), np.linspace(0, 1, 11)) - 1
+        for i in range(10):
+            if (time_bin == i).sum() != 0:
+                terms[f"bin_{i}"] = {"mse": mse_per_instance[time_bin == i].mean()}
+
+        return terms, {}
+    
+    @torch.no_grad()
+    def run_snapshot(
+        self,
+        num_samples: int,
+        batch_size: int,
+        verbose: bool = False,
+    ) -> Dict:
+        # Use current step as seed to ensure different samples for each snapshot
+        import random
+        snapshot_seed = self.step
+        random.seed(snapshot_seed)
+        np.random.seed(snapshot_seed)
+        
+        g = torch.Generator()
+        g.manual_seed(snapshot_seed)
+        
+        dataloader = DataLoader(
+            copy.deepcopy(self.dataset),
+            batch_size=num_samples,
+            shuffle=True,
+            num_workers=0,
+            collate_fn=self.dataset.collate_fn if hasattr(self.dataset, 'collate_fn') else None,
+            generator=g,
+        )
+        data = next(iter(dataloader))
+
+        # Collect metadata (dataset_name and sha256) for wandb display
+        sample_metadata = []
+        if '_dataset_name' in data and '_sha256' in data:
+            for j in range(min(num_samples, len(data['_dataset_name']))):
+                sample_metadata.append(f"{data['_dataset_name'][j]}/{data['_sha256'][j]}")
+        # Remove metadata fields before inference
+        data.pop('_dataset_name', None)
+        data.pop('_sha256', None)
+
+        # inference
+        sampler = self.get_sampler()
+        sample = []
+        cond_vis = []
+        for i in range(0, num_samples, batch_size):
+            batch_data = {k: v[i:i+batch_size] for k, v in data.items()}
+            batch_data = recursive_to_device(batch_data, 'cuda')
+            noise = batch_data['x_0'].replace(torch.randn_like(batch_data['x_0'].feats))
+            cond_vis.append(self.vis_cond(**batch_data))
+            del batch_data['x_0']
+            args = self.get_inference_cond(**batch_data)
+            res = sampler.sample(
+                self.models['denoiser'],
+                noise=noise,
+                **args,
+                steps=12, guidance_strength=3.0, verbose=verbose,
+            )
+            sample.append(res.samples)
+        sample = sp.sparse_cat(sample)
+        
+        sample_gt = {k: v for k, v in data.items()}
+        sample = {k: v if k != 'x_0' else sample for k, v in data.items()}
+        sample_dict = {
+            'sample_gt': {'value': sample_gt, 'type': 'sample'},
+            'sample': {'value': sample, 'type': 'sample'},
+        }
+        if sample_metadata:
+            sample_dict['_metadata'] = sample_metadata
+        sample_dict.update(dict_reduce(cond_vis, None, {
+            'value': lambda x: torch.cat(x, dim=0),
+            'type': lambda x: x[0],
+        }))
+        
+        return sample_dict
+
+
+class SparseFlowMatchingCFGTrainer(ClassifierFreeGuidanceMixin, SparseFlowMatchingTrainer):
+    """
+    Trainer for sparse diffusion model with flow matching objective and classifier-free guidance.
+    
+    Args:
+        models (dict[str, nn.Module]): Models to train.
+        dataset (torch.utils.data.Dataset): Dataset.
+        output_dir (str): Output directory.
+        load_dir (str): Load directory.
+        step (int): Step to load.
+        batch_size (int): Batch size.
+        batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
+        batch_split (int): Split batch with gradient accumulation.
+        max_steps (int): Max steps.
+        optimizer (dict): Optimizer config.
+        lr_scheduler (dict): Learning rate scheduler config.
+        elastic (dict): Elastic memory management config.
+        grad_clip (float or dict): Gradient clip config.
+        ema_rate (float or list): Exponential moving average rates.
+        fp16_mode (str): FP16 mode.
+            - None: No FP16.
+            - 'inflat_all': Hold a inflated fp32 master param for all params.
+            - 'amp': Automatic mixed precision.
+        fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
+        finetune_ckpt (dict): Finetune checkpoint.
+        log_param_stats (bool): Log parameter stats.
+        i_print (int): Print interval.
+        i_log (int): Log interval.
+        i_sample (int): Sample interval.
+        i_save (int): Save interval.
+        i_ddpcheck (int): DDP check interval.
+
+        t_schedule (dict): Time schedule for flow matching.
+        sigma_min (float): Minimum noise level.
+        p_uncond (float): Probability of dropping conditions.
+    """
+    pass
+
+
+class TextConditionedSparseFlowMatchingCFGTrainer(TextConditionedMixin, SparseFlowMatchingCFGTrainer):
+    """
+    Trainer for sparse text-conditioned diffusion model with flow matching objective and classifier-free guidance.
+    
+    Args:
+        models (dict[str, nn.Module]): Models to train.
+        dataset (torch.utils.data.Dataset): Dataset.
+        output_dir (str): Output directory.
+        load_dir (str): Load directory.
+        step (int): Step to load.
+        batch_size (int): Batch size.
+        batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
+        batch_split (int): Split batch with gradient accumulation.
+        max_steps (int): Max steps.
+        optimizer (dict): Optimizer config.
+        lr_scheduler (dict): Learning rate scheduler config.
+        elastic (dict): Elastic memory management config.
+        grad_clip (float or dict): Gradient clip config.
+        ema_rate (float or list): Exponential moving average rates.
+        fp16_mode (str): FP16 mode.
+            - None: No FP16.
+            - 'inflat_all': Hold a inflated fp32 master param for all params.
+            - 'amp': Automatic mixed precision.
+        fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
+        finetune_ckpt (dict): Finetune checkpoint.
+        log_param_stats (bool): Log parameter stats.
+        i_print (int): Print interval.
+        i_log (int): Log interval.
+        i_sample (int): Sample interval.
+        i_save (int): Save interval.
+        i_ddpcheck (int): DDP check interval.
+
+        t_schedule (dict): Time schedule for flow matching.
+        sigma_min (float): Minimum noise level.
+        p_uncond (float): Probability of dropping conditions.
+        text_cond_model(str): Text conditioning model.
+    """
+    pass
+
+
+class ImageConditionedSparseFlowMatchingCFGTrainer(ImageConditionedMixin, SparseFlowMatchingCFGTrainer):
+    """
+    Trainer for sparse image-conditioned diffusion model with flow matching objective and classifier-free guidance.
+    
+    Args:
+        models (dict[str, nn.Module]): Models to train.
+        dataset (torch.utils.data.Dataset): Dataset.
+        output_dir (str): Output directory.
+        load_dir (str): Load directory.
+        step (int): Step to load.
+        batch_size (int): Batch size.
+        batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
+        batch_split (int): Split batch with gradient accumulation.
+        max_steps (int): Max steps.
+        optimizer (dict): Optimizer config.
+        lr_scheduler (dict): Learning rate scheduler config.
+        elastic (dict): Elastic memory management config.
+        grad_clip (float or dict): Gradient clip config.
+        ema_rate (float or list): Exponential moving average rates.
+        fp16_mode (str): FP16 mode.
+            - None: No FP16.
+            - 'inflat_all': Hold a inflated fp32 master param for all params.
+            - 'amp': Automatic mixed precision.
+        fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
+        finetune_ckpt (dict): Finetune checkpoint.
+        log_param_stats (bool): Log parameter stats.
+        i_print (int): Print interval.
+        i_log (int): Log interval.
+        i_sample (int): Sample interval.
+        i_save (int): Save interval.
+        i_ddpcheck (int): DDP check interval.
+
+        t_schedule (dict): Time schedule for flow matching.
+        sigma_min (float): Minimum noise level.
+        p_uncond (float): Probability of dropping conditions.
+        image_cond_model (str): Image conditioning model.
+    """
+    pass
+
+
+class MultiImageConditionedSparseFlowMatchingCFGTrainer(MultiImageConditionedMixin, SparseFlowMatchingCFGTrainer):
+    """
+    Trainer for sparse image-conditioned diffusion model with flow matching objective and classifier-free guidance.
+    
+    Args:
+        models (dict[str, nn.Module]): Models to train.
+        dataset (torch.utils.data.Dataset): Dataset.
+        output_dir (str): Output directory.
+        load_dir (str): Load directory.
+        step (int): Step to load.
+        batch_size (int): Batch size.
+        batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
+        batch_split (int): Split batch with gradient accumulation.
+        max_steps (int): Max steps.
+        optimizer (dict): Optimizer config.
+        lr_scheduler (dict): Learning rate scheduler config.
+        elastic (dict): Elastic memory management config.
+        grad_clip (float or dict): Gradient clip config.
+        ema_rate (float or list): Exponential moving average rates.
+        fp16_mode (str): FP16 mode.
+            - None: No FP16.
+            - 'inflat_all': Hold a inflated fp32 master param for all params.
+            - 'amp': Automatic mixed precision.
+        fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
+        finetune_ckpt (dict): Finetune checkpoint.
+        log_param_stats (bool): Log parameter stats.
+        i_print (int): Print interval.
+        i_log (int): Log interval.
+        i_sample (int): Sample interval.
+        i_save (int): Save interval.
+        i_ddpcheck (int): DDP check interval.
+
+        t_schedule (dict): Time schedule for flow matching.
+        sigma_min (float): Minimum noise level.
+        p_uncond (float): Probability of dropping conditions.
+        image_cond_model (str): Image conditioning model.
+    """
+    pass
+
+
+class ImageConditionedProjSparseFlowMatchingCFGTrainer(ImageConditionedProjMixin, SparseFlowMatchingCFGTrainer):
+    """
+    Trainer for sparse image-conditioned diffusion model with view-aligned projection.
+    
+    Uses ImageConditionedProjMixin for 3D-to-2D feature projection with camera parameters.
+    CFG dropout is handled by ClassifierFreeGuidanceMixin (via p_uncond parameter).
+    
+    The projection grid outputs a full [B, R, R, R, D] tensor, and this trainer extracts
+    features at sparse coordinates using advanced indexing.
+    
+    Args:
+        t_schedule (dict): Time schedule for flow matching.
+        sigma_min (float): Minimum noise level.
+        p_uncond (float): Probability of dropping conditions.
+        image_cond_model (dict): Image conditioning model config (DinoV3ProjFeatureExtractor).
+        run_projection_test (bool): Whether to run projection visualization test before training.
+    """
+    
+    def __init__(self, *args, run_projection_test: bool = True, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.run_projection_test = run_projection_test
+    
+    def training_losses(
+        self,
+        x_0: sp.SparseTensor,
+        cond=None,
+        **kwargs
+    ) -> Tuple[Dict, Dict]:
+        """
+        Compute training losses for a single timestep.
+        
+        Overridden to pass coords from x_0 to get_cond for sparse feature extraction.
+
+        Args:
+            x_0: The [N x ... x C] sparse tensor of the inputs.
+            cond: The [N x ...] tensor of additional conditions.
+            kwargs: Additional arguments to pass to the backbone.
+
+        Returns:
+            a dict with the key "loss" containing a tensor of shape [N].
+            may also contain other keys for different terms.
+        """
+        noise = x_0.replace(torch.randn_like(x_0.feats))
+        t = self.sample_t(x_0.shape[0]).to(x_0.device).float()
+        x_t = self.diffuse(x_0, t, noise=noise)
+        
+        # Pass coords to get_cond for sparse feature extraction from full grid
+        kwargs['coords'] = x_0.coords
+        cond = self.get_cond(cond, **kwargs)
+        
+        # Pass concat_cond to denoiser if present (needed for PBR/texture training
+        # where shape latent is concatenated with PBR latent as input)
+        denoiser_kwargs = {}
+        if 'concat_cond' in kwargs:
+            denoiser_kwargs['concat_cond'] = kwargs['concat_cond']
+        pred = self.training_models['denoiser'](x_t, t * 1000, cond, **denoiser_kwargs)
+        
+        assert pred.shape == noise.shape == x_0.shape
+        target = self.get_v(x_0, noise, t)
+        terms = edict()
+        terms["mse"] = F.mse_loss(pred.feats, target.feats)
+        terms["loss"] = terms["mse"]
+
+        # log loss with time bins
+        mse_per_instance = np.array([
+            F.mse_loss(pred.feats[x_0.layout[i]], target.feats[x_0.layout[i]]).item()
+            for i in range(x_0.shape[0])
+        ])
+        time_bin = np.digitize(t.cpu().numpy(), np.linspace(0, 1, 11)) - 1
+        for i in range(10):
+            if (time_bin == i).sum() != 0:
+                terms[f"bin_{i}"] = {"mse": mse_per_instance[time_bin == i].mean()}
+
+        return terms, {}
+    
+    @torch.no_grad()
+    def run_snapshot(
+        self,
+        num_samples: int,
+        batch_size: int,
+        verbose: bool = False,
+    ) -> Dict:
+        """
+        Run snapshot with coords passed to get_inference_cond for sparse feature extraction.
+        
+        For projection mode, we need to pass coords to properly extract features at
+        sparse positions from the full projection grid.
+        """
+        # Use current step as seed to ensure different samples for each snapshot
+        import random
+        snapshot_seed = self.step
+        random.seed(snapshot_seed)
+        np.random.seed(snapshot_seed)
+        
+        g = torch.Generator()
+        g.manual_seed(snapshot_seed)
+        
+        dataloader = DataLoader(
+            copy.deepcopy(self.dataset),
+            batch_size=num_samples,
+            shuffle=True,
+            num_workers=0,
+            collate_fn=self.dataset.collate_fn if hasattr(self.dataset, 'collate_fn') else None,
+            generator=g,
+        )
+        data = next(iter(dataloader))
+
+        # Collect metadata (dataset_name and sha256) for wandb display
+        sample_metadata = []
+        if '_dataset_name' in data and '_sha256' in data:
+            for j in range(min(num_samples, len(data['_dataset_name']))):
+                sample_metadata.append(f"{data['_dataset_name'][j]}/{data['_sha256'][j]}")
+        # Remove metadata fields before inference
+        data.pop('_dataset_name', None)
+        data.pop('_sha256', None)
+
+        # inference
+        sampler = self.get_sampler()
+        sample = []
+        cond_vis = []
+        for i in range(0, num_samples, batch_size):
+            batch_data = {k: v[i:i+batch_size] for k, v in data.items()}
+            batch_data = recursive_to_device(batch_data, 'cuda')
+            noise = batch_data['x_0'].replace(torch.randn_like(batch_data['x_0'].feats))
+            cond_vis.append(self.vis_cond(**batch_data))
+            
+            # Save coords before deleting x_0 (needed for projection feature extraction)
+            coords = batch_data['x_0'].coords
+            del batch_data['x_0']
+            
+            # Pass coords to get_inference_cond for sparse feature extraction
+            batch_data['coords'] = coords
+            args = self.get_inference_cond(**batch_data)
+            
+            res = sampler.sample(
+                self.models['denoiser'],
+                noise=noise,
+                **args,
+                steps=12, guidance_strength=3.0, verbose=verbose,
+            )
+            sample.append(res.samples)
+        sample = sp.sparse_cat(sample)
+        
+        sample_gt = {k: v for k, v in data.items()}
+        sample = {k: v if k != 'x_0' else sample for k, v in data.items()}
+        sample_dict = {
+            'sample_gt': {'value': sample_gt, 'type': 'sample'},
+            'sample': {'value': sample, 'type': 'sample'},
+        }
+        if sample_metadata:
+            sample_dict['_metadata'] = sample_metadata
+        sample_dict.update(dict_reduce(cond_vis, None, {
+            'value': lambda x: torch.cat(x, dim=0),
+            'type': lambda x: x[0],
+        }))
+        
+        return sample_dict
+    
+    @torch.no_grad()
+    def visualize_sample(self, sample):
+        """
+        Convert a sample to images, including GT camera view if available.
+        
+        Args:
+            sample: Either a SparseTensor or dict containing:
+                - 'x_0': SparseTensor
+                - 'camera_angle_x': [B] (optional)
+                - 'camera_distance': [B] (optional)
+                - 'mesh_scale': [B] (optional)
+                
+        Returns:
+            dict with visualization images or tensor
+        """
+        if hasattr(self.dataset, 'visualize_sample'):
+            if isinstance(sample, dict):
+                # Extract camera params and pass them explicitly, since some
+                # dataset.visualize_sample() (e.g. SLatShapeVisMixin) expect
+                # separate keyword arguments rather than a single dict.
+                camera_kwargs = {}
+                for k in ('camera_angle_x', 'camera_distance', 'mesh_scale'):
+                    if k in sample:
+                        camera_kwargs[k] = sample[k]
+                
+                # Try passing camera kwargs explicitly first; fall back to
+                # passing the entire dict if the dataset method doesn't accept them
+                # (e.g. SLatPbrVisMixin expects a dict with 'x_0' + 'concat_cond').
+                import inspect
+                sig = inspect.signature(self.dataset.visualize_sample)
+                params = list(sig.parameters.keys())
+                if 'camera_angle_x' in params:
+                    # Shape-style: visualize_sample(x_0, camera_angle_x=, ...)
+                    x_0 = sample.get('x_0', sample)
+                    return self.dataset.visualize_sample(x_0, **camera_kwargs)
+                else:
+                    # Tex/PBR-style: visualize_sample(sample_dict)
+                    return self.dataset.visualize_sample(sample)
+            else:
+                return self.dataset.visualize_sample(sample)
+        else:
+            if isinstance(sample, dict):
+                return sample.get('x_0', sample)
+            return sample
+    
+    def run(self):
+        """
+        Run training with projection visualization test before starting.
+        """
+        # Run projection visualization test before training starts (if enabled)
+        if self.run_projection_test and self.is_master:
+            print('\n' + '='*60)
+            print('Running projection visualization test...')
+            print('='*60)
+            self._run_projection_visualization_test()
+            
+        super().run()
+    
+    @torch.no_grad()
+    def _run_projection_visualization_test(self, num_samples: int = 4):
+        """
+        Run projection visualization test on a few samples before training starts.
+        
+        This helps verify that the 3D-to-2D projection is working correctly.
+        """
+        import os
+        from torch.utils.data import DataLoader
+        
+        # Create a small dataloader
+        dataloader = DataLoader(
+            self.dataset,
+            batch_size=min(num_samples, self.snapshot_batch_size),
+            shuffle=True,
+            num_workers=0,
+            collate_fn=self.dataset.collate_fn if hasattr(self.dataset, 'collate_fn') else None,
+        )
+        
+        # Get one batch
+        data = next(iter(dataloader))
+        data = {k: v.cuda() if isinstance(v, torch.Tensor) else v for k, v in data.items()}
+        
+        # Extract condition image
+        cond = data.get('cond')
+        if cond is None:
+            print("Warning: No 'cond' field in data, skipping projection visualization test")
+            return
+        
+        # Save directory
+        save_dir = os.path.join(self.output_dir, 'samples', 'projection_test')
+        
+        # Call visualization method
+        if hasattr(self, 'visualize_projection_test'):
+            # Need to pass camera info as kwargs
+            kwargs = {k: v for k, v in data.items() if k != 'cond' and k != 'x_0'}
+            self.visualize_projection_test(
+                cond=cond,
+                save_dir=save_dir,
+                prefix="proj_test",
+                **kwargs
+            )
+            print(f"Projection visualization saved to: {save_dir}")
+        else:
+            print("Warning: visualize_projection_test not available")
diff --git a/pixal3d/trainers/utils.py b/pixal3d/trainers/utils.py
new file mode 100644
index 0000000..23e4286
--- /dev/null
+++ b/pixal3d/trainers/utils.py
@@ -0,0 +1,91 @@
+import torch
+import torch.nn as nn
+
+# FP16 utils
+from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
+
+
+def str_to_dtype(dtype_str: str):
+    return {
+        'f16': torch.float16,
+        'fp16': torch.float16,
+        'float16': torch.float16,
+        'bf16': torch.bfloat16,
+        'bfloat16': torch.bfloat16,
+        'f32': torch.float32,
+        'fp32': torch.float32,
+        'float32': torch.float32,
+    }[dtype_str]
+    
+
+def make_master_params(model_params):
+    """
+    Copy model parameters into a inflated tensor of full-precision parameters.
+    """
+    master_params = _flatten_dense_tensors(
+        [param.detach().float() for param in model_params]
+    )
+    master_params = nn.Parameter(master_params)
+    master_params.requires_grad = True
+    return [master_params]
+
+
+def unflatten_master_params(model_params, master_params):
+    """
+    Unflatten the master parameters to look like model_params.
+    """
+    return _unflatten_dense_tensors(master_params[0].detach(), model_params)
+
+
+def model_params_to_master_params(model_params, master_params):
+    """
+    Copy the model parameter data into the master parameters.
+    """
+    master_params[0].detach().copy_(
+        _flatten_dense_tensors([param.detach().float() for param in model_params])
+    )
+
+
+def master_params_to_model_params(model_params, master_params):
+    """
+    Copy the master parameter data back into the model parameters.
+    """
+    for param, master_param in zip(
+        model_params, _unflatten_dense_tensors(master_params[0].detach(), model_params)
+    ):
+        param.detach().copy_(master_param)
+
+
+def model_grads_to_master_grads(model_params, master_params):
+    """
+    Copy the gradients from the model parameters into the master parameters
+    from make_master_params().
+    """
+    master_params[0].grad = _flatten_dense_tensors(
+        [param.grad.data.detach().float() for param in model_params]
+    )
+    
+
+def zero_grad(model_params):
+    for param in model_params:
+       if param.grad is not None:
+            if param.grad.grad_fn is not None:
+                param.grad.detach_()
+            else:
+                param.grad.requires_grad_(False)
+            param.grad.zero_()
+            
+
+# LR Schedulers
+from torch.optim.lr_scheduler import LambdaLR
+
+class LinearWarmupLRScheduler(LambdaLR):
+    def __init__(self, optimizer, warmup_steps, last_epoch=-1):
+        self.warmup_steps = warmup_steps
+        super(LinearWarmupLRScheduler, self).__init__(optimizer, self.lr_lambda, last_epoch=last_epoch)
+        
+    def lr_lambda(self, current_step):
+        if current_step < self.warmup_steps:
+            return float(current_step + 1) / self.warmup_steps
+        return 1.0
+        
\ No newline at end of file
diff --git a/pixal3d/trainers/vae/pbr_vae.py b/pixal3d/trainers/vae/pbr_vae.py
new file mode 100644
index 0000000..28d3646
--- /dev/null
+++ b/pixal3d/trainers/vae/pbr_vae.py
@@ -0,0 +1,291 @@
+from typing import *
+import os
+import copy
+import functools
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+import utils3d
+from easydict import EasyDict as edict
+
+from ..basic import BasicTrainer
+from ...modules import sparse as sp
+from ...renderers import MeshRenderer
+from ...representations import Mesh, MeshWithPbrMaterial, MeshWithVoxel
+from ...utils.data_utils import recursive_to_device, cycle, BalancedResumableSampler
+from ...utils.loss_utils import l1_loss, l2_loss, ssim, lpips
+
+
+class PbrVaeTrainer(BasicTrainer):
+    """
+    Trainer for PBR attributes VAE
+    
+    Args:
+        models (dict[str, nn.Module]): Models to train.
+        dataset (torch.utils.data.Dataset): Dataset.
+        output_dir (str): Output directory.
+        load_dir (str): Load directory.
+        step (int): Step to load.
+        batch_size (int): Batch size.
+        batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
+        batch_split (int): Split batch with gradient accumulation.
+        max_steps (int): Max steps.
+        optimizer (dict): Optimizer config.
+        lr_scheduler (dict): Learning rate scheduler config.
+        elastic (dict): Elastic memory management config.
+        grad_clip (float or dict): Gradient clip config.
+        ema_rate (float or list): Exponential moving average rates.
+        fp16_mode (str): FP16 mode.
+            - None: No FP16.
+            - 'inflat_all': Hold a inflated fp32 master param for all params.
+            - 'amp': Automatic mixed precision.
+        fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
+        finetune_ckpt (dict): Finetune checkpoint.
+        log_param_stats (bool): Log parameter stats.
+        i_print (int): Print interval.
+        i_log (int): Log interval.
+        i_sample (int): Sample interval.
+        i_save (int): Save interval.
+        i_ddpcheck (int): DDP check interval.
+        
+        loss_type (str): Loss type.
+        lambda_kl (float): KL loss weight.
+        lambda_ssim (float): SSIM loss weight.
+        lambda_lpips (float): LPIPS loss weight.
+    """
+    
+    def __init__(
+        self,
+        *args,
+        loss_type: str = 'l1',
+        lambda_kl: float = 1e-6,
+        lambda_ssim: float = 0.2,
+        lambda_lpips: float = 0.2,
+        lambda_render: float = 1.0,
+        render_resolution: float = 1024,
+        camera_randomization_config: dict = {
+            'radius_range': [2, 100],
+        },
+        **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        self.loss_type = loss_type
+        self.lambda_kl = lambda_kl
+        self.lambda_ssim = lambda_ssim
+        self.lambda_lpips = lambda_lpips
+        self.lambda_render = lambda_render
+        self.camera_randomization_config = camera_randomization_config
+        
+        self.renderer = MeshRenderer({'near': 1, 'far': 3, 'resolution': render_resolution}, device=self.device)
+        
+    def prepare_dataloader(self, **kwargs):
+        """
+        Prepare dataloader.
+        """
+        self.data_sampler = BalancedResumableSampler(
+            self.dataset,
+            shuffle=True,
+            batch_size=self.batch_size_per_gpu,
+        )
+        self.dataloader = DataLoader(
+            self.dataset,
+            batch_size=self.batch_size_per_gpu,
+            num_workers=int(np.ceil(os.cpu_count() / torch.cuda.device_count())),
+            pin_memory=True,
+            drop_last=True,
+            persistent_workers=True,
+            collate_fn=functools.partial(self.dataset.collate_fn, split_size=self.batch_split),
+            sampler=self.data_sampler,
+        )
+        self.data_iterator = cycle(self.dataloader)
+        
+    def _randomize_camera(self, num_samples: int):
+        # sample radius and fov
+        r_min, r_max = self.camera_randomization_config['radius_range']
+        k_min = 1 / r_max**2
+        k_max = 1 / r_min**2
+        ks = torch.rand(num_samples, device=self.device) * (k_max - k_min) + k_min
+        radius = 1 / torch.sqrt(ks)
+        fov = 2 * torch.arcsin(0.5 / radius)
+        origin = radius.unsqueeze(-1) * F.normalize(torch.randn(num_samples, 3, device=self.device), dim=-1)
+        
+        # build camera
+        extrinsics = utils3d.torch.extrinsics_look_at(origin, torch.zeros_like(origin), torch.tensor([0, 0, 1], dtype=torch.float32, device=self.device))
+        intrinsics = utils3d.torch.intrinsics_from_fov_xy(fov, fov)
+        near = [np.random.uniform(r - 1, r) for r in radius.tolist()]
+        
+        return {
+            'extrinsics': extrinsics,
+            'intrinsics': intrinsics,
+            'near': near,
+        }
+        
+    def _render_batch(self, reps: List[Mesh], extrinsics: torch.Tensor, intrinsics: torch.Tensor, near: List,
+        ) -> Dict[str, torch.Tensor]:
+        """
+        Render a batch of representations.
+
+        Args:
+            reps: The dictionary of lists of representations.
+            extrinsics: The [N x 4 x 4] tensor of extrinsics.
+            intrinsics: The [N x 3 x 3] tensor of intrinsics.
+            
+        Returns: 
+            a dict with
+                base_color : [N x 3 x H x W] tensor of base color.
+                metallic : [N x 1 x H x W] tensor of metallic.
+                roughness : [N x 1 x H x W] tensor of roughness.
+                alpha : [N x 1 x H x W] tensor of alpha.
+        """
+        ret = {k : [] for k in ['base_color', 'metallic', 'roughness', 'alpha']}
+        for i, rep in enumerate(reps):
+            self.renderer.rendering_options['near'] = near[i]
+            self.renderer.rendering_options['far'] = near[i] + 2
+            out_dict = self.renderer.render(rep, extrinsics[i], intrinsics[i], return_types=['attr'])
+            for k in out_dict:
+                ret[k].append(out_dict[k])
+        for k in ret:
+            ret[k] = torch.stack(ret[k])
+        return ret
+    
+    def training_losses(
+        self,
+        x: sp.SparseTensor,
+        mesh: List[MeshWithPbrMaterial] = None,
+        **kwargs
+    ) -> Tuple[Dict, Dict]:
+        """
+        Compute training losses.
+
+        Args:
+            x (SparseTensor): Input sparse tensor for pbr materials.
+            mesh (List[MeshWithPbrMaterial]): The list of meshes with PBR materials.
+
+        Returns:
+            a dict with the key "loss" containing a scalar tensor.
+            may also contain other keys for different terms.
+
+        """
+        z, mean, logvar = self.training_models['encoder'](x, sample_posterior=True, return_raw=True)
+        y = self.training_models['decoder'](z)
+        
+        terms = edict(loss = 0.0)
+        
+        # direct regression
+        if self.loss_type == 'l1':
+            terms["l1"] = l1_loss(x.feats, y.feats)
+            terms["loss"] = terms["loss"] + terms["l1"]
+        elif self.loss_type == 'l2':
+            terms["l2"] = l2_loss(x.feats, y.feats)
+            terms["loss"] = terms["loss"] + terms["l2"]
+        else:
+            raise ValueError(f'Invalid loss type {self.loss_type}')
+        
+        # rendering loss
+        if self.lambda_render != 0.0:
+            recon = [MeshWithVoxel(
+                m.vertices,
+                m.faces,
+                [-0.5, -0.5, -0.5],
+                1 / self.dataset.resolution,
+                v.coords[:, 1:],
+                v.feats * 0.5 + 0.5,
+                torch.Size([*v.shape, *v.spatial_shape]),
+                layout={
+                    'base_color': slice(0, 3),
+                    'metallic': slice(3, 4),
+                    'roughness': slice(4, 5),
+                    'alpha': slice(5, 6),
+                }
+            ) for m, v in zip(mesh, y)]
+            cameras = self._randomize_camera(len(mesh))
+            gt_renders = self._render_batch(mesh, **cameras)
+            pred_renders = self._render_batch(recon, **cameras)
+            gt_base_color = gt_renders['base_color']
+            pred_base_color = pred_renders['base_color']
+            gt_mra = torch.cat([gt_renders['metallic'], gt_renders['roughness'], gt_renders['alpha']], dim=1)
+            pred_mra = torch.cat([pred_renders['metallic'], pred_renders['roughness'], pred_renders['alpha']], dim=1)
+            terms['render/base_color/ssim'] = 1 - ssim(pred_base_color, gt_base_color)
+            terms['render/base_color/lpips'] = lpips(pred_base_color, gt_base_color)
+            terms['render/mra/ssim'] = 1 - ssim(pred_mra, gt_mra)
+            terms['render/mra/lpips'] = lpips(pred_mra, gt_mra)
+            terms['loss'] = terms['loss'] + \
+                            self.lambda_render * (self.lambda_ssim * terms['render/base_color/ssim'] + self.lambda_lpips * terms['render/base_color/lpips'] + \
+                                                self.lambda_ssim * terms['render/mra/ssim'] + self.lambda_lpips * terms['render/mra/lpips'])
+            
+        # KL regularization
+        terms["kl"] = 0.5 * torch.mean(mean.pow(2) + logvar.exp() - logvar - 1)
+        terms["loss"] = terms["loss"] + self.lambda_kl * terms["kl"]
+            
+        return terms, {}
+    
+    @torch.no_grad()
+    def run_snapshot(
+        self,
+        num_samples: int,
+        batch_size: int,
+        verbose: bool = False,
+    ) -> Dict:
+        # Use current step as seed to ensure different samples for each snapshot
+        import random
+        snapshot_seed = self.step
+        random.seed(snapshot_seed)
+        np.random.seed(snapshot_seed)
+        
+        g = torch.Generator()
+        g.manual_seed(snapshot_seed)
+        
+        dataloader = DataLoader(
+            copy.deepcopy(self.dataset),
+            batch_size=batch_size,
+            shuffle=True,
+            num_workers=1,
+            collate_fn=self.dataset.collate_fn if hasattr(self.dataset, 'collate_fn') else None,
+            generator=g,
+        )
+        dataloader.dataset.with_mesh = True
+
+        # inference
+        gts = []
+        recons = []
+        self.models['encoder'].eval()
+        self.models['decoder'].eval()
+        for i in range(0, num_samples, batch_size):
+            batch = min(batch_size, num_samples - i)
+            data = next(iter(dataloader))
+            args = {k: v[:batch] for k, v in data.items()}
+            args = recursive_to_device(args, self.device)
+            z = self.models['encoder'](args['x'])
+            y = self.models['decoder'](z)
+            gts.extend(args['mesh'])
+            recons.extend([MeshWithVoxel(
+                m.vertices,
+                m.faces,
+                [-0.5, -0.5, -0.5],
+                1 / self.dataset.resolution,
+                v.coords[:, 1:],
+                v.feats * 0.5 + 0.5,
+                torch.Size([*v.shape, *v.spatial_shape]),
+                layout={
+                    'base_color': slice(0, 3),
+                    'metallic': slice(3, 4),
+                    'roughness': slice(4, 5),
+                    'alpha': slice(5, 6),
+                }
+            ) for m, v in zip(args['mesh'], y)])
+        self.models['encoder'].train()
+        self.models['decoder'].train()
+        
+        cameras = self._randomize_camera(num_samples)
+        gt_renders = self._render_batch(gts, **cameras)
+        pred_renders = self._render_batch(recons, **cameras)
+
+        sample_dict = {
+            'gt_base_color': {'value': gt_renders['base_color'] * 2 - 1, 'type': 'image'},
+            'pred_base_color': {'value': pred_renders['base_color'] * 2 - 1, 'type': 'image'},
+            'gt_mra': {'value': torch.cat([gt_renders['metallic'], gt_renders['roughness'], gt_renders['alpha']], dim=1) * 2 - 1, 'type': 'image'},
+            'pred_mra': {'value': torch.cat([pred_renders['metallic'], pred_renders['roughness'], pred_renders['alpha']], dim=1) * 2 - 1, 'type': 'image'},
+        }
+            
+        return sample_dict
diff --git a/pixal3d/trainers/vae/shape_vae.py b/pixal3d/trainers/vae/shape_vae.py
new file mode 100644
index 0000000..db4630e
--- /dev/null
+++ b/pixal3d/trainers/vae/shape_vae.py
@@ -0,0 +1,276 @@
+from typing import *
+import os
+import copy
+import functools
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+import utils3d
+from easydict import EasyDict as edict
+
+from ..basic import BasicTrainer
+from ...modules import sparse as sp
+from ...renderers import MeshRenderer
+from ...representations import Mesh
+from ...utils.data_utils import recursive_to_device, cycle, BalancedResumableSampler
+from ...utils.loss_utils import l1_loss, ssim, lpips
+
+
+class ShapeVaeTrainer(BasicTrainer):
+    """
+    Trainer for Shape VAE
+    
+    Args:
+        models (dict[str, nn.Module]): Models to train.
+        dataset (torch.utils.data.Dataset): Dataset.
+        output_dir (str): Output directory.
+        load_dir (str): Load directory.
+        step (int): Step to load.
+        batch_size (int): Batch size.
+        batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
+        batch_split (int): Split batch with gradient accumulation.
+        max_steps (int): Max steps.
+        optimizer (dict): Optimizer config.
+        lr_scheduler (dict): Learning rate scheduler config.
+        elastic (dict): Elastic memory management config.
+        grad_clip (float or dict): Gradient clip config.
+        ema_rate (float or list): Exponential moving average rates.
+        fp16_mode (str): FP16 mode.
+            - None: No FP16.
+            - 'inflat_all': Hold a inflated fp32 master param for all params.
+            - 'amp': Automatic mixed precision.
+        fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
+        finetune_ckpt (dict): Finetune checkpoint.
+        log_param_stats (bool): Log parameter stats.
+        i_print (int): Print interval.
+        i_log (int): Log interval.
+        i_sample (int): Sample interval.
+        i_save (int): Save interval.
+        i_ddpcheck (int): DDP check interval.
+        
+        lambda_subdiv (float): Subdivision loss weight.
+        lambda_intersected (float): Intersected loss weight.
+        lambda_vertice (float): Vertice loss weight.
+        lambda_kl (float): KL loss weight.
+        lambda_ssim (float): SSIM loss weight.
+        lambda_lpips (float): LPIPS loss weight.
+    """
+    
+    def __init__(
+        self,
+        *args,
+        lambda_subdiv: float = 0.1,
+        lambda_intersected: float = 0.1,
+        lambda_vertice: float = 1e-2,
+        lambda_mask: float = 1,
+        lambda_depth: float = 10,
+        lambda_normal: float = 1,
+        lambda_kl: float = 1e-6,
+        lambda_ssim: float = 0.2,
+        lambda_lpips: float = 0.2,
+        render_resolution: float = 1024,
+        camera_randomization_config: dict = {
+            'radius_range': [2, 100],
+        },
+        **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        self.lambda_subdiv = lambda_subdiv
+        self.lambda_intersected = lambda_intersected
+        self.lambda_mask = lambda_mask
+        self.lambda_vertice = lambda_vertice
+        self.lambda_depth = lambda_depth
+        self.lambda_normal = lambda_normal
+        self.lambda_kl = lambda_kl
+        self.lambda_ssim = lambda_ssim
+        self.lambda_lpips = lambda_lpips
+        self.camera_randomization_config = camera_randomization_config
+        
+        self.renderer = MeshRenderer({'near': 1, 'far': 3, 'resolution': render_resolution}, device=self.device)
+        
+    def prepare_dataloader(self, **kwargs):
+        """
+        Prepare dataloader.
+        """
+        self.data_sampler = BalancedResumableSampler(
+            self.dataset,
+            shuffle=True,
+            batch_size=self.batch_size_per_gpu,
+        )
+        self.dataloader = DataLoader(
+            self.dataset,
+            batch_size=self.batch_size_per_gpu,
+            num_workers=int(np.ceil(os.cpu_count() / torch.cuda.device_count())),
+            pin_memory=True,
+            drop_last=True,
+            persistent_workers=True,
+            collate_fn=functools.partial(self.dataset.collate_fn, split_size=self.batch_split),
+            sampler=self.data_sampler,
+        )
+        self.data_iterator = cycle(self.dataloader)
+        
+    def _randomize_camera(self, num_samples: int):
+        # sample radius and fov
+        r_min, r_max = self.camera_randomization_config['radius_range']
+        k_min = 1 / r_max**2
+        k_max = 1 / r_min**2
+        ks = torch.rand(num_samples, device=self.device) * (k_max - k_min) + k_min
+        radius = 1 / torch.sqrt(ks)
+        fov = 2 * torch.arcsin(0.5 / radius)
+        origin = radius.unsqueeze(-1) * F.normalize(torch.randn(num_samples, 3, device=self.device), dim=-1)
+        
+        # build camera
+        extrinsics = utils3d.torch.extrinsics_look_at(origin, torch.zeros_like(origin), torch.tensor([0, 0, 1], dtype=torch.float32, device=self.device))
+        intrinsics = utils3d.torch.intrinsics_from_fov_xy(fov, fov)
+        near = [np.random.uniform(r - 1, r) for r in radius.tolist()]
+        
+        return {
+            'extrinsics': extrinsics,
+            'intrinsics': intrinsics,
+            'near': near,
+        }
+        
+    def _render_batch(self, reps: List[Mesh], extrinsics: torch.Tensor, intrinsics: torch.Tensor, near: List,
+                      return_types=['mask', 'normal', 'depth']) -> Dict[str, torch.Tensor]:
+        """
+        Render a batch of representations.
+
+        Args:
+            reps: The dictionary of lists of representations.
+            extrinsics: The [N x 4 x 4] tensor of extrinsics.
+            intrinsics: The [N x 3 x 3] tensor of intrinsics.
+            return_types: vary in ['mask', 'normal', 'depth', 'normal_map', 'color']
+            
+        Returns: 
+            a dict with
+                mask : [N x 1 x H x W] tensor of rendered masks
+                normal : [N x 3 x H x W] tensor of rendered normals
+                depth : [N x 1 x H x W] tensor of rendered depths
+        """
+        ret = {k : [] for k in return_types}
+        for i, rep in enumerate(reps):
+            self.renderer.rendering_options['near'] = near[i]
+            self.renderer.rendering_options['far'] = near[i] + 2
+            out_dict = self.renderer.render(rep, extrinsics[i], intrinsics[i], return_types=return_types)
+            for k in out_dict:
+                ret[k].append(out_dict[k][None] if k in ['mask', 'depth'] else out_dict[k])
+        for k in ret:
+            ret[k] = torch.stack(ret[k])
+        return ret
+    
+    def training_losses(
+        self,
+        vertices: sp.SparseTensor,
+        intersected: sp.SparseTensor,
+        mesh: List[Mesh],
+    ) -> Tuple[Dict, Dict]:
+        """
+        Compute training losses.
+
+        Args:
+            vertices (SparseTensor): vertices of each active voxel
+            intersected (SparseTensor): intersected flag of each active voxel
+            mesh (List[Mesh]): the list of meshes to render
+
+        Returns:
+            a dict with the key "loss" containing a scalar tensor.
+            may also contain other keys for different terms.
+        """
+        z, mean, logvar = self.training_models['encoder'](vertices, intersected, sample_posterior=True, return_raw=True)
+        recon, pred_vertice, pred_intersected, subs_gt, subs = self.training_models['decoder'](z, intersected)
+        
+        terms = edict(loss = 0.0)
+        
+        # direct regression
+        if self.lambda_intersected > 0:
+            terms["direct/intersected"] = F.binary_cross_entropy_with_logits(pred_intersected.feats.flatten(), intersected.feats.flatten().float())
+            terms["loss"] = terms["loss"] + self.lambda_intersected * terms["direct/intersected"]
+        if self.lambda_vertice > 0:
+            terms["direct/vertice"] = F.mse_loss(pred_vertice.feats, vertices.feats)
+            terms["loss"] = terms["loss"] + self.lambda_vertice * terms["direct/vertice"]
+            
+        # subdivision prediction loss
+        for i, (sub_gt, sub) in enumerate(zip(subs_gt, subs)):
+            terms[f"bce_sub{i}"] = F.binary_cross_entropy_with_logits(sub.feats, sub_gt.float())
+            terms["loss"] = terms["loss"] + self.lambda_subdiv * terms[f"bce_sub{i}"]
+            
+        # rendering loss
+        cameras = self._randomize_camera(len(mesh))
+        gt_renders = self._render_batch(mesh, **cameras, return_types=['mask', 'normal', 'depth'])
+        pred_renders = self._render_batch(recon, **cameras, return_types=['mask', 'normal', 'depth'])
+        terms['render/mask'] = l1_loss(pred_renders['mask'], gt_renders['mask'])
+        terms['render/depth'] = l1_loss(pred_renders['depth'], gt_renders['depth'])
+        terms['render/normal/l1'] = l1_loss(pred_renders['normal'], gt_renders['normal'])
+        terms['render/normal/ssim'] = 1 - ssim(pred_renders['normal'], gt_renders['normal'])
+        terms['render/normal/lpips'] = lpips(pred_renders['normal'], gt_renders['normal'])
+        terms['loss'] = terms['loss'] + \
+                        self.lambda_mask * terms['render/mask'] + \
+                        self.lambda_depth * terms['render/depth'] + \
+                        self.lambda_normal * (terms['render/normal/l1'] + self.lambda_ssim * terms['render/normal/ssim'] + self.lambda_lpips * terms['render/normal/lpips'])
+       
+        # KL regularization     
+        terms["kl"] = 0.5 * torch.mean(mean.pow(2) + logvar.exp() - logvar - 1)
+        terms["loss"] = terms["loss"] + self.lambda_kl * terms["kl"]
+            
+        return terms, {}
+    
+    @torch.no_grad()
+    def run_snapshot(
+        self,
+        num_samples: int,
+        batch_size: int,
+        verbose: bool = False,
+    ) -> Dict:
+        # Use current step as seed to ensure different samples for each snapshot
+        import random
+        snapshot_seed = self.step
+        random.seed(snapshot_seed)
+        np.random.seed(snapshot_seed)
+        
+        g = torch.Generator()
+        g.manual_seed(snapshot_seed)
+        
+        dataloader = DataLoader(
+            copy.deepcopy(self.dataset),
+            batch_size=batch_size,
+            shuffle=True,
+            num_workers=1,
+            collate_fn=self.dataset.collate_fn if hasattr(self.dataset, 'collate_fn') else None,
+            generator=g,
+        )
+
+        # inference
+        gts = []
+        recons = []
+        recons2 = []
+        self.models['encoder'].eval()
+        for i in range(0, num_samples, batch_size):
+            batch = min(batch_size, num_samples - i)
+            data = next(iter(dataloader))
+            args = {k: v[:batch] for k, v in data.items()}
+            args = recursive_to_device(args, self.device)
+            z = self.models['encoder'](args['vertices'], args['intersected'])
+            self.models['decoder'].train()
+            y = self.models['decoder'](z, args['intersected'])[0]
+            z.clear_spatial_cache()
+            self.models['decoder'].eval()
+            y2 = self.models['decoder'](z)
+            gts.extend(args['mesh'])
+            recons.extend(y)
+            recons2.extend(y2)
+        self.models['encoder'].train()
+        self.models['decoder'].train()
+        
+        cameras = self._randomize_camera(num_samples)
+        gt_renders = self._render_batch(gts, **cameras, return_types=['normal'])
+        recons_renders = self._render_batch(recons, **cameras, return_types=['normal'])
+        recons2_renders = self._render_batch(recons2, **cameras, return_types=['normal'])
+        
+        sample_dict = {
+            'gt': {'value': gt_renders['normal'], 'type': 'image'},
+            'rec': {'value': recons_renders['normal'], 'type': 'image'},
+            'rec2': {'value': recons2_renders['normal'], 'type': 'image'},
+        }
+            
+        return sample_dict
diff --git a/pixal3d/trainers/vae/sparse_structure_vae.py b/pixal3d/trainers/vae/sparse_structure_vae.py
new file mode 100644
index 0000000..cbd9d7b
--- /dev/null
+++ b/pixal3d/trainers/vae/sparse_structure_vae.py
@@ -0,0 +1,140 @@
+from typing import *
+import copy
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from easydict import EasyDict as edict
+
+from ..basic import BasicTrainer
+
+
+class SparseStructureVaeTrainer(BasicTrainer):
+    """
+    Trainer for Sparse Structure VAE.
+    
+    Args:
+        models (dict[str, nn.Module]): Models to train.
+        dataset (torch.utils.data.Dataset): Dataset.
+        output_dir (str): Output directory.
+        load_dir (str): Load directory.
+        step (int): Step to load.
+        batch_size (int): Batch size.
+        batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
+        batch_split (int): Split batch with gradient accumulation.
+        max_steps (int): Max steps.
+        optimizer (dict): Optimizer config.
+        lr_scheduler (dict): Learning rate scheduler config.
+        elastic (dict): Elastic memory management config.
+        grad_clip (float or dict): Gradient clip config.
+        ema_rate (float or list): Exponential moving average rates.
+        fp16_mode (str): FP16 mode.
+            - None: No FP16.
+            - 'inflat_all': Hold a inflated fp32 master param for all params.
+            - 'amp': Automatic mixed precision.
+        fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
+        finetune_ckpt (dict): Finetune checkpoint.
+        log_param_stats (bool): Log parameter stats.
+        i_print (int): Print interval.
+        i_log (int): Log interval.
+        i_sample (int): Sample interval.
+        i_save (int): Save interval.
+        i_ddpcheck (int): DDP check interval.
+        
+        loss_type (str): Loss type. 'bce' for binary cross entropy, 'l1' for L1 loss, 'dice' for Dice loss.
+        lambda_kl (float): KL divergence loss weight.
+    """
+    
+    def __init__(
+        self,
+        *args,
+        loss_type='bce',
+        lambda_kl=1e-6,
+        **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        self.loss_type = loss_type
+        self.lambda_kl = lambda_kl
+    
+    def training_losses(
+        self,
+        ss: torch.Tensor,
+        **kwargs
+    ) -> Tuple[Dict, Dict]:
+        """
+        Compute training losses.
+
+        Args:
+            ss: The [N x 1 x H x W x D] tensor of binary sparse structure.
+
+        Returns:
+            a dict with the key "loss" containing a scalar tensor.
+            may also contain other keys for different terms.
+        """
+        z, mean, logvar = self.training_models['encoder'](ss.float(), sample_posterior=True, return_raw=True)
+        logits = self.training_models['decoder'](z)
+
+        terms = edict(loss = 0.0)
+        if self.loss_type == 'bce':
+            terms["bce"] = F.binary_cross_entropy_with_logits(logits, ss.float(), reduction='mean')
+            terms["loss"] = terms["loss"] + terms["bce"]
+        elif self.loss_type == 'l1':
+            terms["l1"] = F.l1_loss(F.sigmoid(logits), ss.float(), reduction='mean')
+            terms["loss"] = terms["loss"] + terms["l1"]
+        elif self.loss_type == 'dice':
+            logits = F.sigmoid(logits)
+            terms["dice"] = 1 - (2 * (logits * ss.float()).sum() + 1) / (logits.sum() + ss.float().sum() + 1)
+            terms["loss"] = terms["loss"] + terms["dice"]
+        else:
+            raise ValueError(f'Invalid loss type {self.loss_type}')
+        terms["kl"] = 0.5 * torch.mean(mean.pow(2) + logvar.exp() - logvar - 1)
+        terms["loss"] = terms["loss"] + self.lamda_kl * terms["kl"]
+            
+        return terms, {}
+    
+    @torch.no_grad()
+    def snapshot(self, suffix=None, num_samples=64, batch_size=1, verbose=False):
+        super().snapshot(suffix=suffix, num_samples=num_samples, batch_size=batch_size, verbose=verbose)
+    
+    @torch.no_grad()
+    def run_snapshot(
+        self,
+        num_samples: int,
+        batch_size: int,
+        verbose: bool = False,
+    ) -> Dict:
+        # Use current step as seed to ensure different samples for each snapshot
+        import random
+        snapshot_seed = self.step
+        random.seed(snapshot_seed)
+        np.random.seed(snapshot_seed)
+        
+        g = torch.Generator()
+        g.manual_seed(snapshot_seed)
+        
+        dataloader = DataLoader(
+            copy.deepcopy(self.dataset),
+            batch_size=batch_size,
+            shuffle=True,
+            num_workers=0,
+            collate_fn=self.dataset.collate_fn if hasattr(self.dataset, 'collate_fn') else None,
+            generator=g,
+        )
+
+        # inference
+        gts = []
+        recons = []
+        for i in range(0, num_samples, batch_size):
+            batch = min(batch_size, num_samples - i)
+            data = next(iter(dataloader))
+            args = {k: v[:batch].cuda() if isinstance(v, torch.Tensor) else v[:batch] for k, v in data.items()}
+            z = self.models['encoder'](args['ss'].float(), sample_posterior=False)
+            logits = self.models['decoder'](z)
+            recon = (logits > 0).long()
+            gts.append(args['ss'])
+            recons.append(recon)
+
+        sample_dict = {
+            'gt': {'value': torch.cat(gts, dim=0), 'type': 'sample'},
+            'recon': {'value': torch.cat(recons, dim=0), 'type': 'sample'},
+        }
+        return sample_dict
diff --git a/pixal3d/utils/__init__.py b/pixal3d/utils/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/pixal3d/utils/data_utils.py b/pixal3d/utils/data_utils.py
new file mode 100644
index 0000000..805b6cc
--- /dev/null
+++ b/pixal3d/utils/data_utils.py
@@ -0,0 +1,226 @@
+from typing import *
+import math
+import torch
+import numpy as np
+from torch.utils.data import Sampler, Dataset, DataLoader, DistributedSampler
+import torch.distributed as dist
+
+
+def recursive_to_device(
+    data: Any,
+    device: torch.device,
+    non_blocking: bool = False,
+) -> Any:
+    """
+    Recursively move all tensors in a data structure to a device.
+    """
+    if hasattr(data, "to"):
+        return data.to(device, non_blocking=non_blocking)
+    elif isinstance(data, (list, tuple)):
+        return type(data)(recursive_to_device(d, device, non_blocking) for d in data)
+    elif isinstance(data, dict):
+        return {k: recursive_to_device(v, device, non_blocking) for k, v in data.items()}
+    else:
+        return data
+
+
+def load_balanced_group_indices(
+    load: List[int],
+    num_groups: int,
+    equal_size: bool = False,
+) -> List[List[int]]:
+    """
+    Split indices into groups with balanced load.
+    """
+    if equal_size:
+        group_size = len(load) // num_groups
+    indices = np.argsort(load)[::-1]
+    groups = [[] for _ in range(num_groups)]
+    group_load = np.zeros(num_groups)
+    for idx in indices:
+        min_group_idx = np.argmin(group_load)
+        groups[min_group_idx].append(idx)
+        if equal_size and len(groups[min_group_idx]) == group_size:
+            group_load[min_group_idx] = float('inf')
+        else:
+            group_load[min_group_idx] += load[idx]
+    return groups
+
+
+def cycle(data_loader: DataLoader) -> Iterator:
+    while True:
+        for data in data_loader:
+            if isinstance(data_loader.sampler, ResumableSampler):
+                data_loader.sampler.idx += data_loader.batch_size   # type: ignore[attr-defined]
+            yield data
+        if isinstance(data_loader.sampler, DistributedSampler):
+            data_loader.sampler.epoch += 1
+        if isinstance(data_loader.sampler, ResumableSampler):
+            data_loader.sampler.epoch += 1
+            data_loader.sampler.idx = 0
+        
+
+class ResumableSampler(Sampler):
+    """
+    Distributed sampler that is resumable.
+
+    Args:
+        dataset: Dataset used for sampling.
+        rank (int, optional): Rank of the current process within :attr:`num_replicas`.
+            By default, :attr:`rank` is retrieved from the current distributed
+            group.
+        shuffle (bool, optional): If ``True`` (default), sampler will shuffle the
+            indices.
+        seed (int, optional): random seed used to shuffle the sampler if
+            :attr:`shuffle=True`. This number should be identical across all
+            processes in the distributed group. Default: ``0``.
+        drop_last (bool, optional): if ``True``, then the sampler will drop the
+            tail of the data to make it evenly divisible across the number of
+            replicas. If ``False``, the sampler will add extra indices to make
+            the data evenly divisible across the replicas. Default: ``False``.
+    """
+
+    def __init__(
+        self,
+        dataset: Dataset,
+        shuffle: bool = True,
+        seed: int = 0,
+        drop_last: bool = False,
+    ) -> None:
+        self.dataset = dataset
+        self.epoch = 0
+        self.idx = 0
+        self.drop_last = drop_last
+        self.world_size = dist.get_world_size() if dist.is_initialized() else 1
+        self.rank = dist.get_rank() if dist.is_initialized() else 0
+        # If the dataset length is evenly divisible by # of replicas, then there
+        # is no need to drop any data, since the dataset will be split equally.
+        if self.drop_last and len(self.dataset) % self.world_size != 0:  # type: ignore[arg-type]
+            # Split to nearest available length that is evenly divisible.
+            # This is to ensure each rank receives the same amount of data when
+            # using this Sampler.
+            self.num_samples = math.ceil(
+                (len(self.dataset) - self.world_size) / self.world_size  # type: ignore[arg-type]
+            )
+        else:
+            self.num_samples = math.ceil(len(self.dataset) / self.world_size)  # type: ignore[arg-type]
+        self.total_size = self.num_samples * self.world_size
+        self.shuffle = shuffle
+        self.seed = seed
+
+    def __iter__(self) -> Iterator:
+        if self.shuffle:
+            # deterministically shuffle based on epoch and seed
+            g = torch.Generator()
+            g.manual_seed(self.seed + self.epoch)
+            indices = torch.randperm(len(self.dataset), generator=g).tolist()  # type: ignore[arg-type]
+        else:
+            indices = list(range(len(self.dataset)))  # type: ignore[arg-type]
+
+        if not self.drop_last:
+            # add extra samples to make it evenly divisible
+            padding_size = self.total_size - len(indices)
+            if padding_size <= len(indices):
+                indices += indices[:padding_size]
+            else:
+                indices += (indices * math.ceil(padding_size / len(indices)))[
+                    :padding_size
+                ]
+        else:
+            # remove tail of data to make it evenly divisible.
+            indices = indices[: self.total_size]
+        assert len(indices) == self.total_size
+
+        # subsample
+        indices = indices[self.rank : self.total_size : self.world_size]
+        
+        # resume from previous state
+        indices = indices[self.idx:]
+
+        return iter(indices)
+
+    def __len__(self) -> int:
+        return self.num_samples
+
+    def state_dict(self) -> dict[str, int]:
+        return {
+            'epoch': self.epoch,
+            'idx': self.idx,
+        }
+        
+    def load_state_dict(self, state_dict):
+        self.epoch = state_dict['epoch']
+        self.idx = state_dict['idx']
+        
+
+class BalancedResumableSampler(ResumableSampler):
+    """
+    Distributed sampler that is resumable and balances the load among the processes.
+
+    Args:
+        dataset: Dataset used for sampling.
+        rank (int, optional): Rank of the current process within :attr:`num_replicas`.
+            By default, :attr:`rank` is retrieved from the current distributed
+            group.
+        shuffle (bool, optional): If ``True`` (default), sampler will shuffle the
+            indices.
+        seed (int, optional): random seed used to shuffle the sampler if
+            :attr:`shuffle=True`. This number should be identical across all
+            processes in the distributed group. Default: ``0``.
+        drop_last (bool, optional): if ``True``, then the sampler will drop the
+            tail of the data to make it evenly divisible across the number of
+            replicas. If ``False``, the sampler will add extra indices to make
+            the data evenly divisible across the replicas. Default: ``False``.
+    """
+
+    def __init__(
+        self,
+        dataset: Dataset,
+        shuffle: bool = True,
+        seed: int = 0,
+        drop_last: bool = False,
+        batch_size: int = 1,
+    ) -> None:
+        assert hasattr(dataset, 'loads'), 'Dataset must have "loads" attribute to use BalancedResumableSampler'
+        super().__init__(dataset, shuffle, seed, drop_last)
+        self.batch_size = batch_size
+        self.loads = dataset.loads
+        
+    def __iter__(self) -> Iterator:
+        if self.shuffle:
+            # deterministically shuffle based on epoch and seed
+            g = torch.Generator()
+            g.manual_seed(self.seed + self.epoch)
+            indices = torch.randperm(len(self.dataset), generator=g).tolist()  # type: ignore[arg-type]
+        else:
+            indices = list(range(len(self.dataset)))  # type: ignore[arg-type]
+
+        if not self.drop_last:
+            # add extra samples to make it evenly divisible
+            padding_size = self.total_size - len(indices)
+            if padding_size <= len(indices):
+                indices += indices[:padding_size]
+            else:
+                indices += (indices * math.ceil(padding_size / len(indices)))[
+                    :padding_size
+                ]
+        else:
+            # remove tail of data to make it evenly divisible.
+            indices = indices[: self.total_size]
+        assert len(indices) == self.total_size
+
+        # balance load among processes
+        num_batches = len(indices) // (self.batch_size * self.world_size)
+        balanced_indices = []
+        for i in range(num_batches):
+            start_idx = i * self.batch_size * self.world_size
+            end_idx = (i + 1) * self.batch_size * self.world_size
+            batch_indices = indices[start_idx:end_idx]
+            batch_loads = [self.loads[idx] for idx in batch_indices]
+            groups = load_balanced_group_indices(batch_loads, self.world_size, equal_size=True)
+            balanced_indices.extend([batch_indices[j] for j in groups[self.rank]])
+        
+        # resume from previous state
+        indices = balanced_indices[self.idx:]
+
+        return iter(indices)
diff --git a/pixal3d/utils/dist_utils.py b/pixal3d/utils/dist_utils.py
new file mode 100644
index 0000000..f57e206
--- /dev/null
+++ b/pixal3d/utils/dist_utils.py
@@ -0,0 +1,94 @@
+import os
+import io
+from datetime import timedelta
+from contextlib import contextmanager
+import torch
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+
+def setup_dist(rank, local_rank, world_size, master_addr, master_port):
+    os.environ['MASTER_ADDR'] = master_addr
+    os.environ['MASTER_PORT'] = master_port
+    os.environ['WORLD_SIZE'] = str(world_size)
+    os.environ['RANK'] = str(rank)
+    os.environ['LOCAL_RANK'] = str(local_rank)
+    torch.cuda.set_device(local_rank)
+    dist.init_process_group('nccl', rank=rank, world_size=world_size, timeout=timedelta(hours=10))
+    
+
+def read_file_dist(path):
+    """
+    Read the binary file distributedly.
+    File is only read once by the rank 0 process and broadcasted to other processes.
+
+    Returns:
+        data (io.BytesIO): The binary data read from the file.
+    """
+    if dist.is_initialized() and dist.get_world_size() > 1:
+        # read file
+        size = torch.LongTensor(1).cuda()
+        if dist.get_rank() == 0:
+            with open(path, 'rb') as f:
+                data = f.read()
+            data = torch.ByteTensor(
+                torch.UntypedStorage.from_buffer(data, dtype=torch.uint8)
+            ).cuda()
+            size[0] = data.shape[0]
+        # broadcast size
+        dist.broadcast(size, src=0)
+        if dist.get_rank() != 0:
+            data = torch.ByteTensor(size[0].item()).cuda()
+        # broadcast data
+        dist.broadcast(data, src=0)
+        # convert to io.BytesIO
+        data = data.cpu().numpy().tobytes()
+        data = io.BytesIO(data)
+        return data
+    else:
+        with open(path, 'rb') as f:
+            data = f.read()
+        data = io.BytesIO(data)
+        return data
+    
+
+def unwrap_dist(model):
+    """
+    Unwrap the model from distributed training.
+    """
+    if isinstance(model, DDP):
+        return model.module
+    return model
+
+
+@contextmanager
+def master_first():
+    """
+    A context manager that ensures master process executes first.
+    """
+    if not dist.is_initialized():
+        yield
+    else:
+        if dist.get_rank() == 0:
+            yield
+            dist.barrier()
+        else:
+            dist.barrier()
+            yield
+            
+
+@contextmanager
+def local_master_first():
+    """
+    A context manager that ensures local master process executes first.
+    """
+    if not dist.is_initialized():
+        yield
+    else:
+        if dist.get_rank() % torch.cuda.device_count() == 0:
+            yield
+            dist.barrier()
+        else:
+            dist.barrier()
+            yield
+    
\ No newline at end of file
diff --git a/pixal3d/utils/elastic_utils.py b/pixal3d/utils/elastic_utils.py
new file mode 100644
index 0000000..cba3cf8
--- /dev/null
+++ b/pixal3d/utils/elastic_utils.py
@@ -0,0 +1,228 @@
+from abc import abstractmethod
+from contextlib import contextmanager
+from typing import Tuple
+import torch
+import torch.nn as nn
+import numpy as np
+
+
+class MemoryController:
+    """
+    Base class for memory management during training.
+    """
+    
+    _last_input_size = None
+    _last_mem_ratio = []
+    
+    @contextmanager
+    def record(self):
+        pass
+    
+    def update_run_states(self, input_size=None, mem_ratio=None):
+        if self._last_input_size is None:
+            self._last_input_size = input_size
+        elif self._last_input_size!= input_size:
+            raise ValueError(f'Input size should not change for different ElasticModules.')
+        self._last_mem_ratio.append(mem_ratio)
+    
+    @abstractmethod
+    def get_mem_ratio(self, input_size):
+        pass
+    
+    @abstractmethod
+    def state_dict(self):
+        pass
+    
+    @abstractmethod
+    def log(self):
+        pass
+
+
+class LinearMemoryController(MemoryController):
+    """
+    A simple controller for memory management during training.
+    The memory usage is modeled as a linear function of:
+        - the number of input parameters
+        - the ratio of memory the model use compared to the maximum usage (with no checkpointing)
+    memory_usage = k * input_size * mem_ratio + b
+    The controller keeps track of the memory usage and gives the
+    expected memory ratio to keep the memory usage under a target
+    """
+    def __init__(
+        self,
+        buffer_size=1000,
+        update_every=500,
+        target_ratio=0.8,
+        available_memory=None,
+        max_mem_ratio_start=0.1,
+        params=None,
+        device=None
+    ):
+        self.buffer_size = buffer_size
+        self.update_every = update_every
+        self.target_ratio = target_ratio
+        self.device = device or torch.cuda.current_device()
+        self.available_memory = available_memory or torch.cuda.get_device_properties(self.device).total_memory / 1024**3
+                
+        self._memory = np.zeros(buffer_size, dtype=np.float32)
+        self._input_size = np.zeros(buffer_size, dtype=np.float32)
+        self._mem_ratio = np.zeros(buffer_size, dtype=np.float32)
+        self._buffer_ptr = 0
+        self._buffer_length = 0
+        self._params = tuple(params) if params is not None else (0.0, 0.0)
+        self._max_mem_ratio = max_mem_ratio_start
+        self.step = 0
+
+    def __repr__(self):
+        return f'LinearMemoryController(target_ratio={self.target_ratio}, available_memory={self.available_memory})'
+        
+    def _add_sample(self, memory, input_size, mem_ratio):
+        self._memory[self._buffer_ptr] = memory
+        self._input_size[self._buffer_ptr] = input_size
+        self._mem_ratio[self._buffer_ptr] = mem_ratio
+        self._buffer_ptr = (self._buffer_ptr + 1) % self.buffer_size
+        self._buffer_length = min(self._buffer_length + 1, self.buffer_size)
+            
+    @contextmanager
+    def record(self):
+        torch.cuda.reset_peak_memory_stats(self.device)
+        self._last_input_size = None
+        self._last_mem_ratio = []
+        yield
+        self._last_memory = torch.cuda.max_memory_allocated(self.device) / 1024**3
+        self._last_mem_ratio = sum(self._last_mem_ratio) / len(self._last_mem_ratio)
+        self._add_sample(self._last_memory, self._last_input_size, self._last_mem_ratio)
+        self.step += 1
+        if self.step % self.update_every == 0:
+            self._max_mem_ratio = min(1.0, self._max_mem_ratio + 0.1)
+            self._fit_params()
+            
+    def _fit_params(self):
+        memory_usage = self._memory[:self._buffer_length]
+        input_size = self._input_size[:self._buffer_length]
+        mem_ratio = self._mem_ratio[:self._buffer_length]
+        
+        x = input_size * mem_ratio
+        y = memory_usage
+        k, b = np.polyfit(x, y, 1)
+        self._params = (k, b)
+        # self._visualize()
+        
+    def _visualize(self):
+        import matplotlib.pyplot as plt
+        memory_usage = self._memory[:self._buffer_length]
+        input_size = self._input_size[:self._buffer_length]
+        mem_ratio = self._mem_ratio[:self._buffer_length]
+        k, b = self._params
+        
+        plt.scatter(input_size * mem_ratio, memory_usage, c=mem_ratio, cmap='viridis')
+        x = np.array([0.0, 20000.0])
+        plt.plot(x, k * x + b, c='r')
+        plt.savefig(f'linear_memory_controller_{self.step}.png')
+        plt.cla()
+        
+    def get_mem_ratio(self, input_size):
+        k, b = self._params
+        if k == 0: return np.random.rand() * self._max_mem_ratio
+        pred = (self.available_memory * self.target_ratio - b) / (k * input_size)
+        return min(self._max_mem_ratio, max(0.0, pred))
+    
+    def state_dict(self):
+        return {
+            'params': self._params,
+        }
+        
+    def load_state_dict(self, state_dict):
+        self._params = tuple(state_dict['params'])
+        
+    def log(self):
+        return {
+            'params/k': self._params[0],
+            'params/b': self._params[1],
+            'memory': self._last_memory,
+            'input_size': self._last_input_size,
+            'mem_ratio': self._last_mem_ratio,
+        }
+    
+    
+class ElasticModule(nn.Module):
+    """
+    Module for training with elastic memory management.
+    """
+    def __init__(self):
+        super().__init__()
+        self._memory_controller: MemoryController = None
+        
+    @abstractmethod
+    def _get_input_size(self, *args, **kwargs) -> int:
+        """
+        Get the size of the input data.
+        
+        Returns:
+            int: The size of the input data.
+        """
+        pass
+    
+    @abstractmethod
+    def _forward_with_mem_ratio(self, *args, mem_ratio=0.0, **kwargs) -> Tuple[float, Tuple]:
+        """
+        Forward with a given memory ratio.
+        """
+        pass
+    
+    def register_memory_controller(self, memory_controller: MemoryController):
+        self._memory_controller = memory_controller
+        
+    def forward(self, *args, **kwargs):
+        if self._memory_controller is None or not torch.is_grad_enabled() or not self.training:
+            _, ret = self._forward_with_mem_ratio(*args, **kwargs)
+        else:
+            input_size = self._get_input_size(*args, **kwargs)
+            mem_ratio = self._memory_controller.get_mem_ratio(input_size)
+            mem_ratio, ret = self._forward_with_mem_ratio(*args, mem_ratio=mem_ratio, **kwargs)
+            self._memory_controller.update_run_states(input_size, mem_ratio)
+        return ret
+    
+
+class ElasticModuleMixin:
+    """
+    Mixin for training with elastic memory management.
+    """
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._memory_controller: MemoryController = None
+        
+    @abstractmethod
+    def _get_input_size(self, *args, **kwargs) -> int:
+        """
+        Get the size of the input data.
+        
+        Returns:
+            int: The size of the input data.
+        """
+        pass
+    
+    @abstractmethod
+    @contextmanager
+    def with_mem_ratio(self, mem_ratio=1.0) -> float:
+        """
+        Context manager for training with a reduced memory ratio compared to the full memory usage.
+        
+        Returns:
+            float: The exact memory ratio used during the forward pass.
+        """
+        pass
+    
+    def register_memory_controller(self, memory_controller: MemoryController):
+        self._memory_controller = memory_controller
+        
+    def forward(self, *args, **kwargs):
+        if self._memory_controller is None or not torch.is_grad_enabled() or not self.training:
+            ret = super().forward(*args, **kwargs)
+        else:
+            input_size = self._get_input_size(*args, **kwargs)
+            mem_ratio = self._memory_controller.get_mem_ratio(input_size)
+            with self.with_mem_ratio(mem_ratio) as exact_mem_ratio:
+                ret = super().forward(*args, **kwargs)
+            self._memory_controller.update_run_states(input_size, exact_mem_ratio)
+        return ret
diff --git a/pixal3d/utils/general_utils.py b/pixal3d/utils/general_utils.py
new file mode 100644
index 0000000..589c103
--- /dev/null
+++ b/pixal3d/utils/general_utils.py
@@ -0,0 +1,373 @@
+import re
+import numpy as np
+import cv2
+import torch
+import contextlib
+
+
+# Dictionary utils
+def _dict_merge(dicta, dictb, prefix=''):
+    """
+    Merge two dictionaries.
+    """
+    assert isinstance(dicta, dict), 'input must be a dictionary'
+    assert isinstance(dictb, dict), 'input must be a dictionary'
+    dict_ = {}
+    all_keys = set(dicta.keys()).union(set(dictb.keys()))
+    for key in all_keys:
+        if key in dicta.keys() and key in dictb.keys():
+            if isinstance(dicta[key], dict) and isinstance(dictb[key], dict):
+                dict_[key] = _dict_merge(dicta[key], dictb[key], prefix=f'{prefix}.{key}')
+            else:
+                raise ValueError(f'Duplicate key {prefix}.{key} found in both dictionaries. Types: {type(dicta[key])}, {type(dictb[key])}')
+        elif key in dicta.keys():
+            dict_[key] = dicta[key]
+        else:
+            dict_[key] = dictb[key]
+    return dict_
+
+
+def dict_merge(dicta, dictb):
+    """
+    Merge two dictionaries.
+    """
+    return _dict_merge(dicta, dictb, prefix='')
+
+
+def dict_foreach(dic, func, special_func={}):
+    """
+    Recursively apply a function to all non-dictionary leaf values in a dictionary.
+    """
+    assert isinstance(dic, dict), 'input must be a dictionary'
+    for key in dic.keys():
+        if isinstance(dic[key], dict):
+            dic[key] = dict_foreach(dic[key], func)
+        else:
+            if key in special_func.keys():
+                dic[key] = special_func[key](dic[key])
+            else:
+                dic[key] = func(dic[key])
+    return dic
+
+
+def dict_reduce(dicts, func, special_func={}):
+    """
+    Reduce a list of dictionaries. Leaf values must be scalars.
+    """
+    assert isinstance(dicts, list), 'input must be a list of dictionaries'
+    assert all([isinstance(d, dict) for d in dicts]), 'input must be a list of dictionaries'
+    assert len(dicts) > 0, 'input must be a non-empty list of dictionaries'
+    all_keys = set([key for dict_ in dicts for key in dict_.keys()])
+    reduced_dict = {}
+    for key in all_keys:
+        vlist = [dict_[key] for dict_ in dicts if key in dict_.keys()]
+        if isinstance(vlist[0], dict):
+            reduced_dict[key] = dict_reduce(vlist, func, special_func)
+        else:
+            if key in special_func.keys():
+                reduced_dict[key] = special_func[key](vlist)
+            else:
+                reduced_dict[key] = func(vlist)
+    return reduced_dict
+
+
+def dict_any(dic, func):
+    """
+    Recursively apply a function to all non-dictionary leaf values in a dictionary.
+    """
+    assert isinstance(dic, dict), 'input must be a dictionary'
+    for key in dic.keys():
+        if isinstance(dic[key], dict):
+            if dict_any(dic[key], func):
+                return True
+        else:
+            if func(dic[key]):
+                return True
+    return False
+
+
+def dict_all(dic, func):
+    """
+    Recursively apply a function to all non-dictionary leaf values in a dictionary.
+    """
+    assert isinstance(dic, dict), 'input must be a dictionary'
+    for key in dic.keys():
+        if isinstance(dic[key], dict):
+            if not dict_all(dic[key], func):
+                return False
+        else:
+            if not func(dic[key]):
+                return False
+    return True
+
+
+def dict_flatten(dic, sep='.'):
+    """
+    Flatten a nested dictionary into a dictionary with no nested dictionaries.
+    """
+    assert isinstance(dic, dict), 'input must be a dictionary'
+    flat_dict = {}
+    for key in dic.keys():
+        if isinstance(dic[key], dict):
+            sub_dict = dict_flatten(dic[key], sep=sep)
+            for sub_key in sub_dict.keys():
+                flat_dict[str(key) + sep + str(sub_key)] = sub_dict[sub_key]
+        else:
+            flat_dict[key] = dic[key]
+    return flat_dict
+
+
+# Context utils
+@contextlib.contextmanager
+def nested_contexts(*contexts):
+    with contextlib.ExitStack() as stack:
+        for ctx in contexts:
+            stack.enter_context(ctx())
+        yield
+
+
+# Image utils
+def make_grid(images, nrow=None, ncol=None, aspect_ratio=None):
+    num_images = len(images)
+    if nrow is None and ncol is None:
+        if aspect_ratio is not None:
+            nrow = int(np.round(np.sqrt(num_images / aspect_ratio)))
+        else:
+            nrow = int(np.sqrt(num_images))
+        ncol = (num_images + nrow - 1) // nrow
+    elif nrow is None and ncol is not None:
+        nrow = (num_images + ncol - 1) // ncol
+    elif nrow is not None and ncol is None:
+        ncol = (num_images + nrow - 1) // nrow
+    else:
+        assert nrow * ncol >= num_images, 'nrow * ncol must be greater than or equal to the number of images'
+    
+    if images[0].ndim == 2:
+        grid = np.zeros((nrow * images[0].shape[0], ncol * images[0].shape[1]), dtype=images[0].dtype)
+    else:
+        grid = np.zeros((nrow * images[0].shape[0], ncol * images[0].shape[1], images[0].shape[2]), dtype=images[0].dtype)
+    for i, img in enumerate(images):
+        row = i // ncol
+        col = i % ncol
+        grid[row * img.shape[0]:(row + 1) * img.shape[0], col * img.shape[1]:(col + 1) * img.shape[1]] = img
+    return grid
+
+
+def notes_on_image(img, notes=None):
+    img = np.pad(img, ((0, 32), (0, 0), (0, 0)), 'constant', constant_values=0)
+    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+    if notes is not None:
+        img = cv2.putText(img, notes, (0, img.shape[0] - 4), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1)
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    return img
+
+
+
+def text_image(text, resolution=(512, 512), max_size=0.5, h_align="left", v_align="center"):
+    """
+    Draw text on an image of the given resolution. The text is automatically wrapped
+    and scaled so that it fits completely within the image while preserving any explicit
+    line breaks and original spacing. Horizontal and vertical alignment can be controlled
+    via flags.
+    
+    Parameters:
+        text (str): The input text. Newline characters and spacing are preserved.
+        resolution (tuple): The image resolution as (width, height).
+        max_size (float): The maximum font size.
+        h_align (str): Horizontal alignment. Options: "left", "center", "right".
+        v_align (str): Vertical alignment. Options: "top", "center", "bottom".
+        
+    Returns:
+        numpy.ndarray: The resulting image (BGR format) with the text drawn.
+    """
+    width, height = resolution
+    # Create a white background image
+    img = np.full((height, width, 3), 255, dtype=np.uint8)
+
+    # Set margins and compute available drawing area
+    margin = 10
+    avail_width = width - 2 * margin
+    avail_height = height - 2 * margin
+
+    # Choose OpenCV font and text thickness
+    font = cv2.FONT_HERSHEY_SIMPLEX
+    thickness = 1
+    # Ratio for additional spacing between lines (relative to the height of "A")
+    line_spacing_ratio = 0.5
+
+    def wrap_line(line, max_width, font, thickness, scale):
+        """
+        Wrap a single line of text into multiple lines such that each line's
+        width (measured at the given scale) does not exceed max_width.
+        This function preserves the original spacing by splitting the line into tokens
+        (words and whitespace) using a regular expression.
+        
+        Parameters:
+            line (str): The input text line.
+            max_width (int): Maximum allowed width in pixels.
+            font (int): OpenCV font identifier.
+            thickness (int): Text thickness.
+            scale (float): The current font scale.
+            
+        Returns:
+            List[str]: A list of wrapped lines.
+        """
+        # Split the line into tokens (words and whitespace), preserving spacing
+        tokens = re.split(r'(\s+)', line)
+        if not tokens:
+            return ['']
+        
+        wrapped_lines = []
+        current_line = ""
+        for token in tokens:
+            candidate = current_line + token
+            candidate_width = cv2.getTextSize(candidate, font, scale, thickness)[0][0]
+            if candidate_width <= max_width:
+                current_line = candidate
+            else:
+                # If current_line is empty, the token itself is too wide;
+                # break the token character by character.
+                if current_line == "":
+                    sub_token = ""
+                    for char in token:
+                        candidate_char = sub_token + char
+                        if cv2.getTextSize(candidate_char, font, scale, thickness)[0][0] <= max_width:
+                            sub_token = candidate_char
+                        else:
+                            if sub_token:
+                                wrapped_lines.append(sub_token)
+                            sub_token = char
+                    current_line = sub_token
+                else:
+                    wrapped_lines.append(current_line)
+                    current_line = token
+        if current_line:
+            wrapped_lines.append(current_line)
+        return wrapped_lines
+
+    def compute_text_block(scale):
+        """
+        Wrap the entire text (splitting at explicit newline characters) using the
+        provided scale, and then compute the overall width and height of the text block.
+        
+        Returns:
+            wrapped_lines (List[str]): The list of wrapped lines.
+            block_width (int): Maximum width among the wrapped lines.
+            block_height (int): Total height of the text block including spacing.
+            sizes (List[tuple]): A list of (width, height) for each wrapped line.
+            spacing (int): The spacing between lines (computed from the scaled "A" height).
+        """
+        # Split text by explicit newlines
+        input_lines = text.splitlines() if text else ['']
+        wrapped_lines = []
+        for line in input_lines:
+            wrapped = wrap_line(line, avail_width, font, thickness, scale)
+            wrapped_lines.extend(wrapped)
+            
+        sizes = []
+        for line in wrapped_lines:
+            (text_size, _) = cv2.getTextSize(line, font, scale, thickness)
+            sizes.append(text_size)  # (width, height)
+            
+        block_width = max((w for w, h in sizes), default=0)
+        # Use the height of "A" (at the current scale) to compute line spacing
+        base_height = cv2.getTextSize("A", font, scale, thickness)[0][1]
+        spacing = int(line_spacing_ratio * base_height)
+        block_height = sum(h for w, h in sizes) + spacing * (len(sizes) - 1) if sizes else 0
+        
+        return wrapped_lines, block_width, block_height, sizes, spacing
+
+    # Use binary search to find the maximum scale that allows the text block to fit
+    lo = 0.001
+    hi = max_size
+    eps = 0.001  # convergence threshold
+    best_scale = lo
+    best_result = None
+
+    while hi - lo > eps:
+        mid = (lo + hi) / 2
+        wrapped_lines, block_width, block_height, sizes, spacing = compute_text_block(mid)
+        # Ensure that both width and height constraints are met
+        if block_width <= avail_width and block_height <= avail_height:
+            best_scale = mid
+            best_result = (wrapped_lines, block_width, block_height, sizes, spacing)
+            lo = mid  # try a larger scale
+        else:
+            hi = mid  # reduce the scale
+
+    if best_result is None:
+        best_scale = 0.5
+        best_result = compute_text_block(best_scale)
+        
+    wrapped_lines, block_width, block_height, sizes, spacing = best_result
+
+    # Compute starting y-coordinate based on vertical alignment flag
+    if v_align == "top":
+        y_top = margin
+    elif v_align == "center":
+        y_top = margin + (avail_height - block_height) // 2
+    elif v_align == "bottom":
+        y_top = margin + (avail_height - block_height)
+    else:
+        y_top = margin + (avail_height - block_height) // 2  # default to center if invalid flag
+
+    # For cv2.putText, the y coordinate represents the text baseline;
+    # so for the first line add its height.
+    y = y_top + (sizes[0][1] if sizes else 0)
+
+    # Draw each line with horizontal alignment based on the flag
+    for i, line in enumerate(wrapped_lines):
+        line_width, line_height = sizes[i]
+        if h_align == "left":
+            x = margin
+        elif h_align == "center":
+            x = margin + (avail_width - line_width) // 2
+        elif h_align == "right":
+            x = margin + (avail_width - line_width)
+        else:
+            x = margin  # default to left if invalid flag
+
+        cv2.putText(img, line, (x, y), font, best_scale, (0, 0, 0), thickness, cv2.LINE_AA)
+        y += line_height + spacing
+
+    return img
+
+
+def save_image_with_notes(img, path, notes=None):
+    """
+    Save an image with notes.
+    """
+    if isinstance(img, torch.Tensor):
+        img = img.cpu().numpy().transpose(1, 2, 0)
+    if img.dtype == np.float32 or img.dtype == np.float64:
+        img = np.clip(img * 255, 0, 255).astype(np.uint8)
+    img = notes_on_image(img, notes)
+    cv2.imwrite(path, cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
+
+
+# debug utils
+
+def atol(x, y):
+    """
+    Absolute tolerance.
+    """
+    return torch.abs(x - y)
+
+
+def rtol(x, y):
+    """
+    Relative tolerance.
+    """
+    return torch.abs(x - y) / torch.clamp_min(torch.maximum(torch.abs(x), torch.abs(y)), 1e-12)
+
+
+# print utils
+def indent(s, n=4):
+    """
+    Indent a string.
+    """
+    lines = s.split('\n')
+    for i in range(1, len(lines)):
+        lines[i] = ' ' * n + lines[i]
+    return '\n'.join(lines)
+
diff --git a/pixal3d/utils/grad_clip_utils.py b/pixal3d/utils/grad_clip_utils.py
new file mode 100644
index 0000000..990a435
--- /dev/null
+++ b/pixal3d/utils/grad_clip_utils.py
@@ -0,0 +1,81 @@
+from typing import *
+import torch
+import numpy as np
+import torch.utils
+
+
+class AdaptiveGradClipper:
+    """
+    Adaptive gradient clipping for training.
+    """
+    def __init__(
+        self,
+        max_norm=None,
+        clip_percentile=95.0,
+        buffer_size=1000,
+    ):
+        self.max_norm = max_norm
+        self.clip_percentile = clip_percentile
+        self.buffer_size = buffer_size
+        
+        self._grad_norm = np.zeros(buffer_size, dtype=np.float32)
+        self._max_norm = max_norm
+        self._buffer_ptr = 0
+        self._buffer_length = 0
+
+    def __repr__(self):
+        return f'AdaptiveGradClipper(max_norm={self.max_norm}, clip_percentile={self.clip_percentile})'
+        
+    def state_dict(self):
+        return {
+            'grad_norm': self._grad_norm,
+            'max_norm': self._max_norm,
+            'buffer_ptr': self._buffer_ptr,
+            'buffer_length': self._buffer_length,
+        }
+
+    def load_state_dict(self, state_dict):
+        self._grad_norm = state_dict['grad_norm']
+        self._max_norm = state_dict['max_norm']
+        self._buffer_ptr = state_dict['buffer_ptr']
+        self._buffer_length = state_dict['buffer_length']
+
+    def log(self):
+        return {
+            'max_norm': self._max_norm,
+        }
+
+    def __call__(self, parameters, norm_type=2.0, error_if_nonfinite=False, foreach=None):
+        """Clip the gradient norm of an iterable of parameters.
+
+        The norm is computed over all gradients together, as if they were
+        concatenated into a single vector. Gradients are modified in-place.
+
+        Args:
+            parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a
+                single Tensor that will have gradients normalized
+            norm_type (float): type of the used p-norm. Can be ``'inf'`` for
+                infinity norm.
+            error_if_nonfinite (bool): if True, an error is thrown if the total
+                norm of the gradients from :attr:`parameters` is ``nan``,
+                ``inf``, or ``-inf``. Default: False (will switch to True in the future)
+            foreach (bool): use the faster foreach-based implementation.
+                If ``None``, use the foreach implementation for CUDA and CPU native tensors and silently
+                fall back to the slow implementation for other device types.
+                Default: ``None``
+
+        Returns:
+            Total norm of the parameter gradients (viewed as a single vector).
+        """
+        max_norm = self._max_norm if self._max_norm is not None else float('inf')
+        grad_norm = torch.nn.utils.clip_grad_norm_(parameters, max_norm=max_norm, norm_type=norm_type, error_if_nonfinite=error_if_nonfinite, foreach=foreach)
+        
+        if torch.isfinite(grad_norm):
+            self._grad_norm[self._buffer_ptr] = grad_norm
+            self._buffer_ptr = (self._buffer_ptr + 1) % self.buffer_size
+            self._buffer_length = min(self._buffer_length + 1, self.buffer_size)
+            if self._buffer_length == self.buffer_size:
+                self._max_norm = np.percentile(self._grad_norm, self.clip_percentile)
+                self._max_norm = min(self._max_norm, self.max_norm) if self.max_norm is not None else self._max_norm
+        
+        return grad_norm
\ No newline at end of file
diff --git a/pixal3d/utils/loss_utils.py b/pixal3d/utils/loss_utils.py
new file mode 100644
index 0000000..52049f6
--- /dev/null
+++ b/pixal3d/utils/loss_utils.py
@@ -0,0 +1,92 @@
+import torch
+import torch.nn.functional as F
+from torch.autograd import Variable
+from math import exp
+from lpips import LPIPS
+
+
+def smooth_l1_loss(pred, target, beta=1.0):
+    diff = torch.abs(pred - target)
+    loss = torch.where(diff < beta, 0.5 * diff ** 2 / beta, diff - 0.5 * beta)
+    return loss.mean()
+
+
+def l1_loss(network_output, gt):
+    return torch.abs((network_output - gt)).mean()
+
+
+def l2_loss(network_output, gt):
+    return ((network_output - gt) ** 2).mean()
+
+
+def gaussian(window_size, sigma):
+    gauss = torch.Tensor([exp(-(x - window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(window_size)])
+    return gauss / gauss.sum()
+
+
+def create_window(window_size, channel):
+    _1D_window = gaussian(window_size, 1.5).unsqueeze(1)
+    _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
+    window = Variable(_2D_window.expand(channel, 1, window_size, window_size).contiguous())
+    return window
+
+
+def psnr(img1, img2, max_val=1.0):
+    mse = F.mse_loss(img1, img2)
+    return 20 * torch.log10(max_val / torch.sqrt(mse))
+
+
+def ssim(img1, img2, window_size=11, size_average=True):
+    channel = img1.size(-3)
+    window = create_window(window_size, channel)
+
+    if img1.is_cuda:
+        window = window.cuda(img1.get_device())
+    window = window.type_as(img1)
+
+    return _ssim(img1, img2, window, window_size, channel, size_average)
+
+def _ssim(img1, img2, window, window_size, channel, size_average=True):
+    mu1 = F.conv2d(img1, window, padding=window_size // 2, groups=channel)
+    mu2 = F.conv2d(img2, window, padding=window_size // 2, groups=channel)
+
+    mu1_sq = mu1.pow(2)
+    mu2_sq = mu2.pow(2)
+    mu1_mu2 = mu1 * mu2
+
+    sigma1_sq = F.conv2d(img1 * img1, window, padding=window_size // 2, groups=channel) - mu1_sq
+    sigma2_sq = F.conv2d(img2 * img2, window, padding=window_size // 2, groups=channel) - mu2_sq
+    sigma12 = F.conv2d(img1 * img2, window, padding=window_size // 2, groups=channel) - mu1_mu2
+
+    C1 = 0.01 ** 2
+    C2 = 0.03 ** 2
+
+    ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))
+
+    if size_average:
+        return ssim_map.mean()
+    else:
+        return ssim_map.mean(1).mean(1).mean(1)
+
+
+loss_fn_vgg = None
+def lpips(img1, img2, value_range=(0, 1)):
+    global loss_fn_vgg
+    if loss_fn_vgg is None:
+        loss_fn_vgg = LPIPS(net='vgg').cuda().eval()
+    # normalize to [-1, 1]
+    img1 = (img1 - value_range[0]) / (value_range[1] - value_range[0]) * 2 - 1
+    img2 = (img2 - value_range[0]) / (value_range[1] - value_range[0]) * 2 - 1
+    return loss_fn_vgg(img1, img2).mean()
+
+
+def normal_angle(pred, gt):
+    pred = pred * 2.0 - 1.0
+    gt = gt * 2.0 - 1.0
+    norms = pred.norm(dim=-1) * gt.norm(dim=-1)
+    cos_sim = (pred * gt).sum(-1) / (norms + 1e-9)
+    cos_sim = torch.clamp(cos_sim, -1.0, 1.0)
+    ang = torch.rad2deg(torch.acos(cos_sim[norms > 1e-9])).mean()
+    if ang.isnan():
+        return -1
+    return ang
diff --git a/pixal3d/utils/mesh_utils.py b/pixal3d/utils/mesh_utils.py
new file mode 100644
index 0000000..a9f1451
--- /dev/null
+++ b/pixal3d/utils/mesh_utils.py
@@ -0,0 +1,268 @@
+from typing import Tuple, Dict
+import numpy as np
+from trimesh import grouping, util, remesh
+import struct
+import re
+from plyfile import PlyData, PlyElement
+
+
+def read_ply(filename):
+    """
+    Read a PLY file and return vertices, triangle faces, and quad faces.
+    
+    Args:
+        filename (str): The file path to read from.
+        
+    Returns:
+        vertices (np.ndarray): Array of shape [N, 3] containing vertex positions.
+        tris (np.ndarray): Array of shape [M, 3] containing triangle face indices (empty if none).
+        quads (np.ndarray): Array of shape [K, 4] containing quad face indices (empty if none).
+    """
+    with open(filename, 'rb') as f:
+        # Read the header until 'end_header' is encountered
+        header_bytes = b""
+        while True:
+            line = f.readline()
+            if not line:
+                raise ValueError("PLY header not found")
+            header_bytes += line
+            if b"end_header" in line:
+                break
+        header = header_bytes.decode('utf-8')
+        
+        # Determine if the file is in ASCII or binary format
+        is_ascii = "ascii" in header
+        
+        # Extract the number of vertices and faces from the header using regex
+        vertex_match = re.search(r'element vertex (\d+)', header)
+        if vertex_match:
+            num_vertices = int(vertex_match.group(1))
+        else:
+            raise ValueError("Vertex count not found in header")
+            
+        face_match = re.search(r'element face (\d+)', header)
+        if face_match:
+            num_faces = int(face_match.group(1))
+        else:
+            raise ValueError("Face count not found in header")
+        
+        vertices = []
+        tris = []
+        quads = []
+        
+        if is_ascii:
+            # For ASCII format, read each line of vertex data (each line contains 3 floats)
+            for _ in range(num_vertices):
+                line = f.readline().decode('utf-8').strip()
+                if not line: 
+                    continue
+                parts = line.split()
+                vertices.append([float(parts[0]), float(parts[1]), float(parts[2])])
+            
+            # Read face data, where the first number indicates the number of vertices for the face
+            for _ in range(num_faces):
+                line = f.readline().decode('utf-8').strip()
+                if not line: 
+                    continue
+                parts = line.split()
+                count = int(parts[0])
+                indices = list(map(int, parts[1:]))
+                if count == 3:
+                    tris.append(indices)
+                elif count == 4:
+                    quads.append(indices)
+                else:
+                    # Skip faces with other numbers of vertices (can be extended as needed)
+                    pass
+        else:
+            # For binary format: read directly from the binary stream
+            # Each vertex consists of 3 floats (12 bytes per vertex)
+            for _ in range(num_vertices):
+                data = f.read(12)
+                if len(data) < 12:
+                    raise ValueError("Insufficient vertex data")
+                v = struct.unpack('<fff', data)
+                vertices.append(v)
+            
+            # Read face data from the binary stream
+            for _ in range(num_faces):
+                # First, read 1 byte indicating the number of vertices in the face
+                count_data = f.read(1)
+                if len(count_data) < 1:
+                    raise ValueError("Failed to read face vertex count")
+                count = struct.unpack('<B', count_data)[0]
+                if count == 3:
+                    data = f.read(12)  # 3 * 4 bytes
+                    if len(data) < 12:
+                        raise ValueError("Insufficient data for triangle face")
+                    indices = struct.unpack('<3i', data)
+                    tris.append(indices)
+                elif count == 4:
+                    data = f.read(16)  # 4 * 4 bytes
+                    if len(data) < 16:
+                        raise ValueError("Insufficient data for quad face")
+                    indices = struct.unpack('<4i', data)
+                    quads.append(indices)
+                else:
+                    # For faces with a different number of vertices, read count*4 bytes
+                    data = f.read(count * 4)
+                    # Skip or extend processing as needed
+                    raise ValueError(f"Unsupported face with {count} vertices")
+        
+        # Convert lists to torch.Tensor
+        vertices = np.array(vertices, dtype=np.float32)
+        tris = np.array(tris, dtype=np.int32) if len(tris) > 0 else np.empty((0, 3), dtype=np.int32)
+        quads = np.array(quads, dtype=np.int32) if len(quads) > 0 else np.empty((0, 4), dtype=np.int32)
+        
+        return vertices, tris, quads
+
+
+def write_ply(
+    filename: str,
+    vertices: np.ndarray,
+    tris: np.ndarray,
+    quads: np.ndarray,
+    vertex_colors: np.ndarray = None,
+    ascii: bool = False
+):
+    """
+    Write a mesh to a PLY file, with the option to save in ASCII or binary format,
+    and optional per-vertex colors.
+    
+    Args:
+        filename (str): The filename to write to.
+        vertices (np.ndarray): [N, 3] The vertex positions.
+        tris (np.ndarray): [M, 3] The triangle indices.
+        quads (np.ndarray): [K, 4] The quad indices.
+        vertex_colors (np.ndarray, optional): [N, 3] or [N, 4] UInt8 colors for each vertex (RGB or RGBA).
+        ascii (bool): If True, write in ASCII format; otherwise binary little-endian.
+    """
+    import struct
+
+    num_vertices = len(vertices)
+    num_faces = len(tris) + len(quads)
+
+    # Build header
+    header_lines = [
+        "ply",
+        f"format {'ascii 1.0' if ascii else 'binary_little_endian 1.0'}",
+        f"element vertex {num_vertices}",
+        "property float x",
+        "property float y",
+        "property float z",
+    ]
+
+    # Add vertex color properties if provided
+    has_color = vertex_colors is not None
+    if has_color:
+        # Expect uint8 values 0-255
+        header_lines += [
+            "property uchar red",
+            "property uchar green",
+            "property uchar blue",
+        ]
+        # Include alpha if RGBA
+        if vertex_colors.shape[1] == 4:
+            header_lines.append("property uchar alpha")
+
+    header_lines += [
+        f"element face {num_faces}",
+        "property list uchar int vertex_index",
+        "end_header",
+        ""
+    ]
+    header = "\n".join(header_lines)
+
+    mode = 'w' if ascii else 'wb'
+    with open(filename, mode) as f:
+        # Write header
+        if ascii:
+            f.write(header)
+        else:
+            f.write(header.encode('utf-8'))
+
+        # Write vertex data
+        for i, v in enumerate(vertices):
+            if ascii:
+                line = f"{v[0]} {v[1]} {v[2]}"
+                if has_color:
+                    col = vertex_colors[i]
+                    line += ' ' + ' '.join(str(int(c)) for c in col)
+                f.write(line + '\n')
+            else:
+                # pack xyz as floats
+                f.write(struct.pack('<fff', *v))
+                if has_color:
+                    col = vertex_colors[i]
+                    # pack as uchar
+                    if col.shape[0] == 3:
+                        f.write(struct.pack('<BBB', *col))
+                    else:
+                        f.write(struct.pack('<BBBB', *col))
+
+        # Write face data
+        if ascii:
+            for tri in tris:
+                f.write(f"3 {tri[0]} {tri[1]} {tri[2]}\n")
+            for quad in quads:
+                f.write(f"4 {quad[0]} {quad[1]} {quad[2]} {quad[3]}\n")
+        else:
+            for tri in tris:
+                f.write(struct.pack('<B3i', 3, *tri))
+            for quad in quads:
+                f.write(struct.pack('<B4i', 4, *quad))
+                
+
+def write_pbr_ply(
+    filename: str,
+    vertices: np.ndarray,
+    faces: np.ndarray,
+    base_color: np.ndarray,
+    metallic: np.ndarray,
+    roughness: np.ndarray,
+    alpha: np.ndarray,
+    ascii: bool = False
+):
+    """
+    Write a mesh to a PLY file, with the option to save in ASCII or binary format,
+    and optional per-vertex colors.
+    
+    Args:
+        filename (str): The filename to write to.
+        vertices (np.ndarray): [N, 3] The vertex positions.
+        faces (np.ndarray): [M, 3] The triangle indices.
+        base_color (np.ndarray): [N, 3] UInt8 colors for each vertex (RGB).
+        metallic (np.ndarray): [N] UInt8 values for metallicness.
+        roughness (np.ndarray): [N] UInt8 values for roughness.
+        alpha (np.ndarray): [N] UInt8 values for alpha.
+        ascii (bool): If True, write in ASCII format; otherwise binary little-endian.
+    """
+    vertex_dtype = [
+        ('x', 'f4'), ('y', 'f4'), ('z', 'f4'),
+        ('red', 'u1'), ('green', 'u1'), ('blue', 'u1'),
+        ('metallic', 'u1'), ('roughness', 'u1'), ('alpha', 'u1')
+    ]
+    
+    vertex_data = np.empty(len(vertices), dtype=vertex_dtype)
+    vertex_data['x'] = vertices[:, 0]
+    vertex_data['y'] = vertices[:, 1]
+    vertex_data['z'] = vertices[:, 2]
+    vertex_data['red'] = base_color[:, 0]
+    vertex_data['green'] = base_color[:, 1]
+    vertex_data['blue'] = base_color[:, 2]
+    vertex_data['metallic'] = metallic
+    vertex_data['roughness'] = roughness
+    vertex_data['alpha'] = alpha
+    
+    face_dtype = [
+        ('vertex_indices', 'i4', (3,))
+    ]
+    
+    face_data = np.empty(len(faces), dtype=face_dtype)
+    face_data['vertex_indices'] = faces
+    
+    ply_data = PlyData([
+        PlyElement.describe(vertex_data,'vertex'),
+        PlyElement.describe(face_data, 'face'),
+    ], text=ascii)
+    ply_data.write(filename)
diff --git a/pixal3d/utils/random_utils.py b/pixal3d/utils/random_utils.py
new file mode 100644
index 0000000..5b668c2
--- /dev/null
+++ b/pixal3d/utils/random_utils.py
@@ -0,0 +1,30 @@
+import numpy as np
+
+PRIMES = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53]
+
+def radical_inverse(base, n):
+    val = 0
+    inv_base = 1.0 / base
+    inv_base_n = inv_base
+    while n > 0:
+        digit = n % base
+        val += digit * inv_base_n
+        n //= base
+        inv_base_n *= inv_base
+    return val
+
+def halton_sequence(dim, n):
+    return [radical_inverse(PRIMES[dim], n) for dim in range(dim)]
+
+def hammersley_sequence(dim, n, num_samples):
+    return [n / num_samples] + halton_sequence(dim - 1, n)
+
+def sphere_hammersley_sequence(n, num_samples, offset=(0, 0), remap=False):
+    u, v = hammersley_sequence(2, n, num_samples)
+    u += offset[0] / num_samples
+    v += offset[1]
+    if remap:
+        u = 2 * u if u < 0.25 else 2 / 3 * u + 1 / 3
+    theta = np.arccos(1 - 2 * u) - np.pi / 2
+    phi = v * 2 * np.pi
+    return [phi, theta]
\ No newline at end of file
diff --git a/pixal3d/utils/render_utils.py b/pixal3d/utils/render_utils.py
new file mode 100644
index 0000000..a3a1165
--- /dev/null
+++ b/pixal3d/utils/render_utils.py
@@ -0,0 +1,231 @@
+import gc
+import torch
+import numpy as np
+from tqdm import tqdm
+import utils3d
+from PIL import Image
+
+from ..renderers import MeshRenderer, VoxelRenderer, PbrMeshRenderer
+from ..representations import Mesh, Voxel, MeshWithPbrMaterial, MeshWithVoxel
+from .random_utils import sphere_hammersley_sequence
+
+
+def yaw_pitch_r_fov_to_extrinsics_intrinsics(yaws, pitchs, rs, fovs):
+    is_list = isinstance(yaws, list)
+    if not is_list:
+        yaws = [yaws]
+        pitchs = [pitchs]
+    if not isinstance(rs, list):
+        rs = [rs] * len(yaws)
+    if not isinstance(fovs, list):
+        fovs = [fovs] * len(yaws)
+    extrinsics = []
+    intrinsics = []
+    for yaw, pitch, r, fov in zip(yaws, pitchs, rs, fovs):
+        fov = torch.deg2rad(torch.tensor(float(fov))).cuda()
+        yaw = torch.tensor(float(yaw)).cuda()
+        pitch = torch.tensor(float(pitch)).cuda()
+        orig = torch.tensor([
+            torch.sin(yaw) * torch.cos(pitch),
+            torch.cos(yaw) * torch.cos(pitch),
+            torch.sin(pitch),
+        ]).cuda() * r
+        extr = utils3d.torch.extrinsics_look_at(orig, torch.tensor([0, 0, 0]).float().cuda(), torch.tensor([0, 0, 1]).float().cuda())
+        intr = utils3d.torch.intrinsics_from_fov_xy(fov, fov)
+        extrinsics.append(extr)
+        intrinsics.append(intr)
+    if not is_list:
+        extrinsics = extrinsics[0]
+        intrinsics = intrinsics[0]
+    return extrinsics, intrinsics
+
+
+def get_renderer(sample, **kwargs):
+    if isinstance(sample, (MeshWithPbrMaterial, MeshWithVoxel)):
+        renderer = PbrMeshRenderer()
+        renderer.rendering_options.resolution = kwargs.get('resolution', 512)
+        renderer.rendering_options.near = kwargs.get('near', 1)
+        renderer.rendering_options.far = kwargs.get('far', 100)
+        renderer.rendering_options.ssaa = kwargs.get('ssaa', 2)
+        renderer.rendering_options.peel_layers = kwargs.get('peel_layers', 8)
+    elif isinstance(sample, Mesh):
+        renderer = MeshRenderer()
+        renderer.rendering_options.resolution = kwargs.get('resolution', 512)
+        renderer.rendering_options.near = kwargs.get('near', 1)
+        renderer.rendering_options.far = kwargs.get('far', 100)
+        renderer.rendering_options.ssaa = kwargs.get('ssaa', 2)
+        renderer.rendering_options.chunk_size = kwargs.get('chunk_size', None)
+    elif isinstance(sample, Voxel):
+        renderer = VoxelRenderer()
+        renderer.rendering_options.resolution = kwargs.get('resolution', 512)
+        renderer.rendering_options.near = kwargs.get('near', 0.1)
+        renderer.rendering_options.far = kwargs.get('far', 10.0)
+        renderer.rendering_options.ssaa = kwargs.get('ssaa', 2)
+    else:
+        raise ValueError(f'Unsupported sample type: {type(sample)}')
+    return renderer
+
+
+@torch.no_grad()
+def render_frames(sample, extrinsics, intrinsics, options={}, verbose=True, renderer=None, **kwargs):
+    if renderer is None:
+        renderer = get_renderer(sample, **options)
+    rets = {}
+    for j, (extr, intr) in tqdm(enumerate(zip(extrinsics, intrinsics)), total=len(extrinsics), desc='Rendering', disable=not verbose):
+        res = renderer.render(sample, extr, intr, **kwargs)
+        for k, v in res.items():
+            if k not in rets: rets[k] = []
+            if v.dim() == 2: v = v[None].repeat(3, 1, 1)
+            rets[k].append(np.clip(v.detach().cpu().numpy().transpose(1, 2, 0) * 255, 0, 255).astype(np.uint8))
+    return rets
+
+
+def render_video(sample, resolution=1024, bg_color=(0, 0, 0), num_frames=40, r=2, fov=40,
+                 start_yaw=None, start_pitch=None, **kwargs):
+    """
+    Render a turntable video of the sample.
+    
+    Args:
+        start_yaw: Starting yaw angle in radians. If None, defaults to π/2.
+        start_pitch: Starting pitch angle in radians. If None, uses the default oscillating pitch
+                     starting at ~0.25.
+    """
+    if start_yaw is None:
+        start_yaw = np.pi / 2
+    yaws = -torch.linspace(0, 2 * 3.1415, num_frames) + start_yaw
+    if start_pitch is not None:
+        pitch = start_pitch + 0.5 * torch.sin(torch.linspace(0, 2 * 3.1415, num_frames))
+    else:
+        pitch = 0.25 + 0.5 * torch.sin(torch.linspace(0, 2 * 3.1415, num_frames))
+    yaws = yaws.tolist()
+    pitch = pitch.tolist()
+    extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaws, pitch, r, fov)
+    return render_frames(sample, extrinsics, intrinsics, {'resolution': resolution, 'bg_color': bg_color}, **kwargs)
+
+
+def render_multiview(sample, resolution=512, nviews=30):
+    r = 2
+    fov = 40
+    cams = [sphere_hammersley_sequence(i, nviews) for i in range(nviews)]
+    yaws = [cam[0] for cam in cams]
+    pitchs = [cam[1] for cam in cams]
+    extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaws, pitchs, r, fov)
+    res = render_frames(sample, extrinsics, intrinsics, {'resolution': resolution, 'bg_color': (0, 0, 0)})
+    return res['color'], extrinsics, intrinsics
+
+
+def render_snapshot(samples, resolution=512, bg_color=(0, 0, 0), offset=(-16 / 180 * np.pi, 20 / 180 * np.pi), r=10, fov=8, nviews=4, **kwargs):
+    yaw = np.linspace(0, 2 * np.pi, nviews, endpoint=False)
+    yaw_offset = offset[0]
+    yaw = [y + yaw_offset for y in yaw]
+    pitch = [offset[1] for _ in range(nviews)]
+    extrinsics, intrinsics = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaw, pitch, r, fov)
+    return render_frames(samples, extrinsics, intrinsics, {'resolution': resolution, 'bg_color': bg_color}, **kwargs)
+
+
+def proj_camera_to_render_params(camera_angle_x, distance):
+    """
+    Convert proj camera parameters to renderer-compatible extrinsics + intrinsics.
+    
+    The proj camera (Blender convention) views from the front. Through empirical
+    testing, this corresponds to extrinsics_look_at from (0, 0, +distance) with up=Y
+    in the mesh coordinate system used by the renderer.
+    
+    Args:
+        camera_angle_x: horizontal FOV in radians
+        distance: camera distance
+    
+    Returns:
+        extrinsics: [4, 4] OpenCV world-to-camera (on CUDA)
+        intrinsics: [3, 3] OpenCV normalized intrinsics (on CUDA)
+    """
+    orig = torch.tensor([0.0, 0.0, distance]).cuda()
+    target = torch.tensor([0.0, 0.0, 0.0]).cuda()
+    up = torch.tensor([0.0, 1.0, 0.0]).cuda()
+    extrinsics = utils3d.torch.extrinsics_look_at(orig, target, up)
+    
+    fov_tensor = torch.tensor(camera_angle_x, dtype=torch.float32).cuda()
+    intrinsics = utils3d.torch.intrinsics_from_fov_xy(fov_tensor, fov_tensor)
+    
+    return extrinsics, intrinsics
+
+
+def render_proj_aligned_video(sample, camera_angle_x, distance, resolution=1024,
+                              num_frames=40, bg_color=(0, 0, 0), **kwargs):
+    """
+    Render a turntable video starting from the proj camera viewpoint.
+    
+    The first frame matches the proj input image exactly. Subsequent frames
+    rotate around the object (around Y axis, which is up in mesh space).
+    
+    Args:
+        sample: mesh to render
+        camera_angle_x: proj camera FOV in radians
+        distance: proj camera distance
+        resolution: render resolution
+        num_frames: number of video frames
+        bg_color: background color
+        **kwargs: additional kwargs (e.g. envmap)
+    
+    Returns:
+        render result dict (same as render_frames)
+    """
+    import math
+    
+    extr_first, intr_first = proj_camera_to_render_params(camera_angle_x, distance)
+    
+    extrinsics_list = []
+    intrinsics_list = []
+    
+    angles = torch.linspace(0, 2 * math.pi, num_frames + 1)[:num_frames]
+    
+    for angle in angles:
+        # Rotation around Y axis (up in mesh space)
+        c = torch.cos(angle)
+        s = torch.sin(angle)
+        R_y = torch.tensor([
+            [ c, 0, s, 0],
+            [ 0, 1, 0, 0],
+            [-s, 0, c, 0],
+            [ 0, 0, 0, 1],
+        ], dtype=torch.float32).cuda()
+        
+        # world-to-camera for rotated world: extr @ R_y^{-1}
+        R_y_inv = R_y.clone()
+        R_y_inv[:3, :3] = R_y[:3, :3].T
+        extr_rotated = extr_first @ R_y_inv
+        
+        extrinsics_list.append(extr_rotated)
+        intrinsics_list.append(intr_first)
+    
+    render_options = {'resolution': resolution, 'bg_color': bg_color}
+    if 'near' in kwargs:
+        render_options['near'] = kwargs.pop('near')
+    if 'far' in kwargs:
+        render_options['far'] = kwargs.pop('far')
+    
+    return render_frames(sample, extrinsics_list, intrinsics_list,
+                         render_options, **kwargs)
+
+
+def make_pbr_vis_frames(result, resolution=1024):
+    num_frames = len(result['shaded'])
+    frames = []
+    for i in range(num_frames):
+        shaded = Image.fromarray(result['shaded'][i])
+        normal = Image.fromarray(result['normal'][i])
+        base_color = Image.fromarray(result['base_color'][i])
+        metallic = Image.fromarray(result['metallic'][i])
+        roughness = Image.fromarray(result['roughness'][i])
+        alpha = Image.fromarray(result['alpha'][i])
+        shaded = shaded.resize((resolution, resolution))
+        normal = normal.resize((resolution, resolution))
+        base_color = base_color.resize((resolution//2, resolution//2))
+        metallic = metallic.resize((resolution//2, resolution//2))
+        roughness = roughness.resize((resolution//2, resolution//2))
+        alpha = alpha.resize((resolution//2, resolution//2))
+        row1 = np.concatenate([shaded, normal], axis=1)
+        row2 = np.concatenate([base_color, metallic, roughness, alpha], axis=1)
+        frame = np.concatenate([row1, row2], axis=0)
+        frames.append(frame)
+    return frames
diff --git a/pixal3d/utils/vis_utils.py b/pixal3d/utils/vis_utils.py
new file mode 100644
index 0000000..0e5f58e
--- /dev/null
+++ b/pixal3d/utils/vis_utils.py
@@ -0,0 +1,44 @@
+from typing import *
+import numpy as np
+import torch
+from ..modules import sparse as sp
+from ..representations import Voxel
+from .render_utils import render_video
+
+
+def pca_color(feats: torch.Tensor, channels: Tuple[int, int, int] = (0, 1, 2)) -> torch.Tensor:
+    """
+    Apply PCA to the features and return the first three principal components.
+    """
+    feats = feats.detach()
+    u, s, v = torch.svd(feats)
+    color = u[:, channels]
+    color = (color - color.min(dim=0, keepdim=True)[0]) / (color.max(dim=0, keepdim=True)[0] - color.min(dim=0, keepdim=True)[0])
+    return color
+    
+
+def vis_sparse_tensor(
+    x: sp.SparseTensor,
+    num_frames: int = 300,
+):
+    assert x.shape[0] == 1, "Only support batch size 1"
+    assert x.coords.shape[1] == 4, "Only support 3D coordinates"
+    
+    coords = x.coords.cuda().detach()[:, 1:]
+    feats = x.feats.cuda().detach()
+    color = pca_color(feats)
+    
+    resolution = max(list(x.spatial_shape))
+    resolution = int(2**np.ceil(np.log2(resolution)))
+    
+    rep = Voxel(
+        origin=[-0.5, -0.5, -0.5],
+        voxel_size=1/resolution,
+        coords=coords,
+        attrs=color,
+        layout={
+            'color': slice(0, 3),
+        }
+    )
+
+    return render_video(rep, colors_overwrite=color, num_frames=num_frames)['color']
diff --git a/requirements-hfdemo.txt b/requirements-hfdemo.txt
new file mode 100644
index 0000000..3588fc5
--- /dev/null
+++ b/requirements-hfdemo.txt
@@ -0,0 +1,28 @@
+--extra-index-url https://download.pytorch.org/whl/cu124
+
+torch==2.6.0
+torchvision==0.21.0
+triton==3.2.0
+pillow==12.0.0
+imageio==2.37.2
+imageio-ffmpeg==0.6.0
+tqdm==4.67.1
+easydict==1.13
+opencv-python-headless==4.12.0.88
+trimesh==4.10.1
+transformers==4.57.3
+zstandard==0.25.0
+kornia==0.8.2
+timm==1.0.22
+diffusers==0.37.1
+accelerate==1.13.0
+gradio
+plyfile==1.1.3
+git+https://github.com/microsoft/MoGe.git
+https://github.com/LDYang694/Storages/releases/download/20260430/natten-0.21.0+torch2.6cu124-cp310-cp310-linux_x86_64.whl
+https://github.com/JeffreyXiang/Storages/releases/download/Space_Wheels_251210/flash_attn_3-3.0.0b1-cp39-abi3-linux_x86_64.whl
+https://github.com/JeffreyXiang/Storages/releases/download/Space_Wheels_251210/cumesh-0.0.1-cp310-cp310-linux_x86_64.whl
+https://github.com/JeffreyXiang/Storages/releases/download/Space_Wheels_251210/flex_gemm-0.0.1-cp310-cp310-linux_x86_64.whl
+https://github.com/JeffreyXiang/Storages/releases/download/Space_Wheels_251210/o_voxel-0.0.1-cp310-cp310-linux_x86_64.whl
+https://github.com/JeffreyXiang/Storages/releases/download/Space_Wheels_251210/nvdiffrast-0.4.0-cp310-cp310-linux_x86_64.whl
+https://github.com/JeffreyXiang/Storages/releases/download/Space_Wheels_251210/nvdiffrec_render-0.0.0-cp310-cp310-linux_x86_64.whl
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..8304c6a
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,16 @@
+git+https://github.com/microsoft/MoGe.git
+pillow==12.0.0
+imageio==2.37.2
+imageio-ffmpeg==0.6.0
+tqdm==4.67.1
+easydict==1.13
+opencv-python-headless==4.12.0.88
+trimesh==4.10.1
+transformers==4.57.3
+zstandard==0.25.0
+kornia==0.8.2
+timm==1.0.22
+diffusers==0.37.1
+accelerate==1.13.0
+gradio
+plyfile==1.1.3
\ No newline at end of file
diff --git a/train.py b/train.py
new file mode 100644
index 0000000..ccd034a
--- /dev/null
+++ b/train.py
@@ -0,0 +1,230 @@
+import os
+import sys
+import json
+import glob
+import argparse
+from easydict import EasyDict as edict
+
+import torch
+import torch.multiprocessing as mp
+import numpy as np
+import random
+
+try:
+    import wandb
+    WANDB_AVAILABLE = True
+except ImportError:
+    WANDB_AVAILABLE = False
+    print("Warning: wandb not installed. Install with 'pip install wandb' to enable wandb logging.")
+
+from pixal3d import models, datasets, trainers
+from pixal3d.utils.dist_utils import setup_dist
+
+
+def find_ckpt(cfg):
+    # Load checkpoint
+    cfg['load_ckpt'] = None
+    if cfg.load_dir != '':
+        if cfg.ckpt == 'latest':
+            files = glob.glob(os.path.join(cfg.load_dir, 'ckpts', 'misc_*.pt'))
+            if len(files) != 0:
+                cfg.load_ckpt = max([
+                    int(os.path.basename(f).split('step')[-1].split('.')[0])
+                    for f in files
+                ])
+        elif cfg.ckpt == 'none':
+            cfg.load_ckpt = None
+        else:
+            cfg.load_ckpt = int(cfg.ckpt)
+    return cfg
+
+
+def setup_rng(rank):
+    torch.manual_seed(rank)
+    torch.cuda.manual_seed_all(rank)
+    np.random.seed(rank)
+    random.seed(rank)
+
+
+def get_model_summary(model):
+    num_params = 0
+    num_trainable_params = 0
+    for name, param in model.named_parameters():
+        num_params += param.numel()
+        if param.requires_grad:
+            num_trainable_params += param.numel()
+    model_summary = f'Number of parameters: {num_params:,}\n'
+    model_summary += f'Number of trainable parameters: {num_trainable_params:,}\n'
+    return model_summary
+
+
+def main(local_rank, cfg):
+    # Set up distributed training
+    rank = cfg.node_rank * cfg.num_gpus + local_rank
+    world_size = cfg.num_nodes * cfg.num_gpus
+    if world_size > 1:
+        setup_dist(rank, local_rank, world_size, cfg.master_addr, cfg.master_port)
+    
+    # Multi-GPU training verification
+    print(f'[Rank {rank}/{world_size}] Process started on GPU {local_rank} (cuda:{torch.cuda.current_device()})')
+    if rank == 0:
+        print(f'\n{"="*60}')
+        print(f'Multi-GPU Training Verification:')
+        print(f'  - Total GPUs (world_size): {world_size}')
+        print(f'  - num_gpus per node: {cfg.num_gpus}')
+        print(f'  - num_nodes: {cfg.num_nodes}')
+        print(f'{"="*60}\n')
+
+    # Initialize wandb (only on rank 0)
+    wandb_run = None
+    if rank == 0 and cfg.use_wandb and WANDB_AVAILABLE:
+        # Use WANDB_DIR env var (local SSD) to avoid S3 FUSE rename/append issues
+        wandb_dir = os.environ.get('WANDB_DIR', cfg.output_dir)
+        os.makedirs(wandb_dir, exist_ok=True)
+        wandb_run = wandb.init(
+            project=cfg.wandb_project,
+            name=cfg.wandb_name if cfg.wandb_name else os.path.basename(cfg.output_dir),
+            config=cfg.__dict__,
+            dir=wandb_dir,
+            resume="allow",
+            id=cfg.wandb_id if cfg.wandb_id else None,
+        )
+        print(f'Wandb initialized: {wandb_run.url}')
+
+        # Upload config JSON file as wandb artifact
+        config_file = cfg.get('config', None)
+        if config_file and os.path.isfile(config_file):
+            config_artifact = wandb.Artifact(
+                name=f"config-{wandb_run.id}",
+                type="config",
+                description=f"Training config for {wandb_run.name}",
+            )
+            config_artifact.add_file(config_file, name=os.path.basename(config_file))
+            # Also save the resolved full config (with command-line args merged)
+            resolved_config_path = os.path.join(cfg.output_dir, 'config_resolved.json')
+            with open(resolved_config_path, 'w') as f:
+                json.dump(cfg.__dict__, f, indent=4, default=str)
+            config_artifact.add_file(resolved_config_path, name='config_resolved.json')
+            wandb_run.log_artifact(config_artifact)
+            print(f'Uploaded config artifact to wandb: {config_file}')
+
+    # Seed rngs
+    setup_rng(rank)
+
+    # Load data
+    dataset_kwargs = dict(cfg.dataset.args)
+    dataset = getattr(datasets, cfg.dataset.name)(cfg.data_dir, **dataset_kwargs)
+
+    # Print dataset info (only on rank 0)
+    if rank == 0:
+        print(f'\nDataset: {cfg.dataset.name}, Number of samples: {len(dataset):,}\n')
+
+    # Build model
+    model_dict = {
+        name: getattr(models, model.name)(**model.args).cuda()
+        for name, model in cfg.models.items()
+    }
+
+    # Model summary
+    if rank == 0:
+        for name, backbone in model_dict.items():
+            model_summary = get_model_summary(backbone)
+            print(f'\n\nBackbone: {name}\n' + model_summary)
+            with open(os.path.join(cfg.output_dir, f'{name}_model_summary.txt'), 'w') as fp:
+                print(model_summary, file=fp)
+
+    # Build trainer
+    trainer = getattr(trainers, cfg.trainer.name)(
+        model_dict, dataset, 
+        **cfg.trainer.args, 
+        output_dir=cfg.output_dir, 
+        load_dir=cfg.load_dir, 
+        step=cfg.load_ckpt,
+        wandb_run=wandb_run,  # Pass wandb run to trainer
+    )
+
+    # Train
+    if not cfg.tryrun:
+        if cfg.profile:
+            trainer.profile()
+        else:
+            trainer.run()
+    
+    # Close wandb
+    if wandb_run is not None:
+        wandb_run.finish()
+
+
+if __name__ == '__main__':
+    # Arguments and config
+    parser = argparse.ArgumentParser()
+    ## config
+    parser.add_argument('--config', type=str, required=True, help='Experiment config file')
+    ## io and resume
+    parser.add_argument('--output_dir', type=str, required=True, help='Output directory')
+    parser.add_argument('--load_dir', type=str, default='', help='Load directory, default to output_dir')
+    parser.add_argument('--ckpt', type=str, default='latest', help='Checkpoint step to resume training, default to latest')
+    parser.add_argument('--data_dir', type=str, default='./data/', help='Data directory')
+    parser.add_argument('--auto_retry', type=int, default=3, help='Number of retries on error')
+    ## dubug
+    parser.add_argument('--tryrun', action='store_true', help='Try run without training')
+    parser.add_argument('--profile', action='store_true', help='Profile training')
+    ## multi-node and multi-gpu
+    parser.add_argument('--num_nodes', type=int, default=1, help='Number of nodes')
+    parser.add_argument('--node_rank', type=int, default=0, help='Node rank')
+    parser.add_argument('--num_gpus', type=int, default=-1, help='Number of GPUs per node, default to all')
+    parser.add_argument('--master_addr', type=str, default='localhost', help='Master address for distributed training')
+    parser.add_argument('--master_port', type=str, default='12666', help='Port for distributed training')
+    ## wandb
+    parser.add_argument('--use_wandb', action='store_true', help='Enable wandb logging')
+    parser.add_argument('--wandb_project', type=str, default='pixal3d-training', help='Wandb project name')
+    parser.add_argument('--wandb_name', type=str, default='', help='Wandb run name, default to output_dir basename')
+    parser.add_argument('--wandb_id', type=str, default='', help='Wandb run id for resuming')
+    opt = parser.parse_args()
+    opt.load_dir = opt.load_dir if opt.load_dir != '' else opt.output_dir
+    opt.num_gpus = torch.cuda.device_count() if opt.num_gpus == -1 else opt.num_gpus
+    ## Load config
+    config = json.load(open(opt.config, 'r'))
+    ## Combine arguments and config
+    cfg = edict()
+    cfg.update(opt.__dict__)
+    cfg.update(config)
+    print('\n\nConfig:')
+    print('=' * 80)
+    print(json.dumps(cfg.__dict__, indent=4))
+
+    # Prepare output directory
+    if cfg.node_rank == 0:
+        os.makedirs(cfg.output_dir, exist_ok=True)
+        ## Save command and config
+        with open(os.path.join(cfg.output_dir, 'command.txt'), 'w') as fp:
+            print(' '.join(['python'] + sys.argv), file=fp)
+        with open(os.path.join(cfg.output_dir, 'config.json'), 'w') as fp:
+            json.dump(config, fp, indent=4)
+
+    # Run
+    if cfg.auto_retry == 0:
+        cfg = find_ckpt(cfg)
+        if cfg.num_gpus > 1:
+            mp.spawn(main, args=(cfg,), nprocs=cfg.num_gpus, join=True)
+        else:
+            main(0, cfg)
+    else:
+        for rty in range(cfg.auto_retry):
+            try:
+                cfg = find_ckpt(cfg)
+                if cfg.num_gpus > 1:
+                    mp.spawn(main, args=(cfg,), nprocs=cfg.num_gpus, join=True)
+                else:
+                    main(0, cfg)
+                break
+            except Exception as e:
+                import traceback
+                print(f'\n{"="*60}')
+                print(f'Error: {e}')
+                print(f'{"="*60}')
+                print('Full traceback:')
+                traceback.print_exc()
+                print(f'{"="*60}')
+                print(f'Retrying ({rty + 1}/{cfg.auto_retry})...')
+            
\ No newline at end of file