Duplicate from wushuang98/Direct3D-S2

Co-authored-by: Shuang Wu <wushuang98@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,36 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/teaserv6.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,109 @@
+---
+license: mit
+tags:
+- image-to-3d
+- 3d-aigc
+- 3d-reconstruction
+- 3d-models
+- 3d-generation
+---
+
+
+# Direct3D‑S2: Gigascale 3D Generation Made Easy with Spatial Sparse Attention
+
+<div align="center">
+  <a href=https://www.neural4d.com/research/direct3d-s2 target="_blank"><img src=https://img.shields.io/badge/Project%20Page-333399.svg?logo=googlehome height=22px></a>
+  <a href=https://huggingface.co/spaces/wushuang98/Direct3D-S2-v1.0-demo target="_blank"><img src=https://img.shields.io/badge/%F0%9F%A4%97%20Demo-276cb4.svg height=22px></a>
+  <a href=https://huggingface.co/spaces/wushuang98/Direct3D-S2-v1.0-demo target="_blank"><img src=https://img.shields.io/badge/%F0%9F%A4%97%20Models-d96902.svg height=22px></a>
+  <a href=https://arxiv.org/pdf/2505.17412 target="_blank"><img src=https://img.shields.io/badge/Arxiv-b5212f.svg?logo=arxiv height=22px></a>
+</div>
+
+<div style="background: #fff; box-shadow: 0 4px 12px rgba(0,0,0,.15); display: inline-block; padding: 0px;">
+    <img id="teaser" src="assets/teaserv6.png" alt="Teaser image of Direct3D-S2"/>
+</div>
+
+---
+
+## ✨ News
+
+- May 30, 2025: 🤯 We have released both v1.0 and v1.1. The new model offers even greater speed compared to FlashAttention-2, with **12.2×** faster forward pass and **19.7×** faster backward pass, resulting in nearly **2×** inference speedup over v1.0.
+- May 30, 2025: 🔨 Release inference code and model.
+- May 26, 2025: 🎁 Release live demo on 🤗 [Hugging Face](https://huggingface.co/spaces/wushuang98/Direct3D-S2-v1.0-demo).
+- May 26, 2025: 🚀 Release paper and project page.
+
+## 📝 Abstract
+
+Generating high-resolution 3D shapes using volumetric representations such as Signed Distance Functions (SDFs) presents substantial computational and memory challenges. We introduce <strong class="has-text-weight-bold">Direct3D‑S2</strong>, a scalable 3D generation framework based on sparse volumes that achieves superior output quality with dramatically reduced training costs. Our key innovation is the <strong class="has-text-weight-bold">Spatial Sparse Attention (SSA)</strong> mechanism, which greatly enhances the efficiency of Diffusion Transformer (DiT) computations on sparse volumetric data. SSA allows the model to effectively process large token sets within sparse volumes, substantially reducing computational overhead and achieving a <em>3.9&times;</em> speedup in the forward pass and a <em>9.6&times;</em> speedup in the backward pass. Our framework also includes a variational autoencoder (VAE) that maintains a consistent sparse volumetric format across input, latent, and output stages. Compared to previous methods with heterogeneous representations in 3D VAE, this unified design significantly improves training efficiency and stability. Our model is trained on public available datasets, and experiments demonstrate that <strong class="has-text-weight-bold">Direct3D‑S2</strong> not only surpasses state-of-the-art methods in generation quality and efficiency, but also enables <strong class="has-text-weight-bold">training at 1024<sup>3</sup>  resolution with just 8 GPUs</strong>, a task typically requiring at least 32 GPUs for volumetric representations at 256<sup>3</sup> resolution, thus making gigascale 3D generation both practical and accessible.
+
+## 🌟 Highlight
+
+- **Gigascale 3D Generation**: Direct3D-S2 enables training at 1024<sup>3</sup> resolution with only 8 GPUs.
+- **Spatial Sparse Attention (SSA)**: A novel attention mechanism designed for sparse volumetric data, enabling efficient processing of large token sets.
+- **Unified Sparse VAE**: A variational autoencoder that maintains a consistent sparse volumetric format across input, latent, and output stages, improving training efficiency and stability.
+
+## 🚀 Getting Started
+
+### Installation
+
+```sh
+git clone https://github.com/DreamTechAI/Direct3D-S2.git
+
+cd Direct3D-S2
+
+pip install -r requirements.txt
+
+pip install -e .
+
+```
+
+### Usage
+
+```python
+from direct3d_s2.pipeline import Direct3DS2Pipeline
+pipeline = Direct3DS2Pipeline.from_pretrained(
+  'wushuang98/Direct3D-S2', 
+  subfolder="direct3d-s2-v-1-1"
+).to("cuda:0")
+
+mesh = pipeline(
+  'assets/test/13.png', 
+  sdf_resolution=1024, # 512 or 1024
+  remesh=False, # Switch to True if you need to reduce the number of triangles.
+)["mesh"]
+
+mesh.export('output.obj')
+```
+
+### Web Demo
+
+We provide a Gradio web demo for Direct3D-S2, which allows you to generate 3D meshes from images interactively.
+
+```bash
+python app.py
+```
+
+## 🤗 Acknowledgements
+
+Thanks to the following repos for their great work, which helps us a lot in the development of Direct3D-S2:
+
+- [Trellis](https://github.com/microsoft/TRELLIS)
+- [SparseFlex](https://github.com/VAST-AI-Research/TripoSF)
+- [native-sparse-attention-triton](https://github.com/XunhaoLai/native-sparse-attention-triton)
+- [diffusers](https://github.com/huggingface/diffusers)
+
+## 📄 License
+
+Direct3D-S2 is released under the MIT License. See [LICENSE](LICENSE) for details.
+
+## 📖 Citation
+
+If you find our work useful, please consider citing our paper:
+
+```bibtex
+@article{wu2025direct3ds2gigascale3dgeneration,
+  title={Direct3D-S2: Gigascale 3D Generation Made Easy with Spatial Sparse Attention}, 
+  author={Shuang Wu and Youtian Lin and Feihu Zhang and Yifei Zeng and Yikang Yang and Yajie Bao and Jiachen Qian and Siyu Zhu and Philip Torr and Xun Cao and Yao Yao},
+  journal={arXiv preprint arXiv:2505.17412},
+  year={2025}
+}
+```

config.json

@@ -0,0 +1,5 @@
+{
+  "Name": [
+    "Direct3D-S2"
+  ],
+}

direct3d-s2-v-1-0/config.yaml

@@ -0,0 +1,171 @@
+dense_vae:
+  target: direct3d_s2.models.autoencoders.dense_vae.DenseShapeVAE
+  params:
+    use_checkpoint: true
+    embed_dim: 8
+    in_channels: 1
+    out_channels: 1
+    model_channels_encoder: [32, 128, 512]  
+    num_res_blocks_encoder: 2
+    num_res_blocks_middle_encoder: 2
+    model_channels_decoder: [512, 128, 64]
+    num_res_blocks_decoder: 4
+    num_res_blocks_middle_decoder: 4
+    use_fp16: true
+    latents_scale: 1.0
+    latents_shift: 0.0
+      
+dense_dit:
+  target: direct3d_s2.models.transformers.dense_dit.DenseDiT
+  params:
+    resolution: 16 
+    in_channels: 8
+    out_channels: 8
+    model_channels: 1024
+    cond_channels: 1024
+    num_blocks: 24
+    num_heads: 16
+    mlp_ratio: 4
+    patch_size: 1
+    pe_mode: ape
+    qk_rms_norm: true
+    use_checkpoint: true
+    use_fp16: true
+    latent_shape: [8, 16, 16, 16]
+
+dense_image_encoder:
+  target: direct3d_s2.models.conditioner.DinoEncoder
+  params:
+    model: facebookresearch/dinov2
+    version: dinov2_vitl14_reg
+    size: 518
+
+dense_scheduler:
+  target: diffusers.schedulers.FlowMatchEulerDiscreteScheduler
+  params:
+    num_train_timesteps: 1000
+    shift: 6.0 
+
+sparse_vae_512:
+  target: direct3d_s2.models.autoencoders.ss_vae.SparseSDFVAE
+  params:
+    use_checkpoint: true
+    embed_dim: 16
+    num_head_channels_encoder: 64
+    model_channels_encoder: 512 
+    num_heads_encoder: 8 
+    num_blocks_encoder: 4
+    num_head_channels_decoder: 64
+    model_channels_decoder: 512 
+    num_heads_decoder: 8 
+    num_blocks_decoder: 4
+    resolution: 64
+    out_channels: 1
+    use_fp16: true
+    latents_scale: 1.0
+    latents_shift: 0.0
+
+sparse_dit_512:
+  target: direct3d_s2.models.transformers.sparse_dit.SparseDiT 
+  params:
+    resolution: 64 
+    in_channels: 16
+    out_channels: 16
+    model_channels: 1024
+    cond_channels: 1024
+    num_blocks: 24
+    num_heads: 32
+    num_kv_heads: 2
+    compression_block_size: 4
+    selection_block_size: 8
+    topk: 32
+    compression_version: v2
+    pe_mode: ape
+    factor: 1.0
+    sparse_conditions: true
+    qk_rms_norm: true
+    use_shift: true
+    use_checkpoint: true
+    use_fp16: true
+
+sparse_scheduler_512:
+  target: diffusers.schedulers.FlowMatchEulerDiscreteScheduler
+  params:
+    num_train_timesteps: 1000
+    shift: 6.0
+
+sparse_vae_1024:
+  target: direct3d_s2.models.autoencoders.ss_vae.SparseSDFVAE
+  params:
+    use_checkpoint: true
+    embed_dim: 16
+    num_head_channels_encoder: 64
+    model_channels_encoder: 512 
+    num_heads_encoder: 8 
+    num_blocks_encoder: 4
+    num_head_channels_decoder: 64
+    model_channels_decoder: 512 
+    num_heads_decoder: 8 
+    num_blocks_decoder: 4
+    resolution: 128
+    out_channels: 1
+    use_fp16: true
+    latents_scale: 1.0
+    latents_shift: 0.0
+    chunk_size: 4
+
+sparse_dit_1024:
+  target: direct3d_s2.models.transformers.sparse_dit.SparseDiT 
+  params:
+    resolution: 128 
+    in_channels: 16
+    out_channels: 16
+    model_channels: 1024
+    cond_channels: 1024
+    num_blocks: 24
+    num_heads: 32
+    num_kv_heads: 2
+    compression_block_size: 4
+    selection_block_size: 8
+    topk: 8
+    compression_version: v2
+    pe_mode: ape
+    factor: 0.5
+    sparse_conditions: true
+    qk_rms_norm: true
+    use_shift: true
+    use_checkpoint: true
+    use_fp16: true
+
+sparse_scheduler_1024:
+  target: diffusers.schedulers.FlowMatchEulerDiscreteScheduler
+  params:
+    num_train_timesteps: 1000
+    shift: 8.0
+
+sparse_image_encoder:
+  target: direct3d_s2.models.conditioner.DinoEncoder
+  params:
+    model: facebookresearch/dinov2
+    version: dinov2_vitl14_reg
+    size: 518
+
+refiner:
+  target: direct3d_s2.models.refiner.unet_refiner.Voxel_RefinerXL
+  params:
+    in_channels: 1
+    out_channels: 1
+    layers_per_block:  2
+    layers_mid_block:  2
+    patch_size: 192
+    use_fp16: true
+
+refiner_1024:
+  target: direct3d_s2.models.refiner.unet_refiner.Voxel_RefinerXL_sign
+  params:
+    in_channels: 1
+    out_channels: 1
+    layers_per_block:  2
+    layers_mid_block:  2
+    patch_size: 256
+    use_fp16: true

direct3d-s2-v-1-0/model_dense.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5e6598155f907ca15eb9e4fb96905f9df4d10ebc319f9567ea947db8c684d8fd
+size 1505767932

direct3d-s2-v-1-0/model_refiner.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2eb0f70d323a43d63ab37c91790c8a7c2b9ada1d7771ffe99ef7a647bdbc99e4
+size 269562302

direct3d-s2-v-1-0/model_refiner_1024.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e75e405067f5c3eed48b2e46ee752e1f6593bf9a0db421cb4d694a677958460a
+size 269498584

direct3d-s2-v-1-0/model_sparse_1024.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b2fe3f7daab2a07b99248665501ddb52ced65b632557688227921b79faf8b9a1
+size 1116239946

direct3d-s2-v-1-0/model_sparse_512.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9cc7fe7d4978cf8ff96121ece226b5344c543486283f05d108b7acb3126f0c69
+size 1116238952

direct3d-s2-v-1-1/config.yaml

@@ -0,0 +1,171 @@
+dense_vae:
+  target: direct3d_s2.models.autoencoders.dense_vae.DenseShapeVAE
+  params:
+    use_checkpoint: true
+    embed_dim: 8
+    in_channels: 1
+    out_channels: 1
+    model_channels_encoder: [32, 128, 512]  
+    num_res_blocks_encoder: 2
+    num_res_blocks_middle_encoder: 2
+    model_channels_decoder: [512, 128, 64]
+    num_res_blocks_decoder: 4
+    num_res_blocks_middle_decoder: 4
+    use_fp16: true
+    latents_scale: 1.0
+    latents_shift: 0.0
+      
+dense_dit:
+  target: direct3d_s2.models.transformers.dense_dit.DenseDiT
+  params:
+    resolution: 16 
+    in_channels: 8
+    out_channels: 8
+    model_channels: 1024
+    cond_channels: 1024
+    num_blocks: 24
+    num_heads: 16
+    mlp_ratio: 4
+    patch_size: 1
+    pe_mode: ape
+    qk_rms_norm: true
+    use_checkpoint: true
+    use_fp16: true
+    latent_shape: [8, 16, 16, 16]
+
+dense_image_encoder:
+  target: direct3d_s2.models.conditioner.DinoEncoder
+  params:
+    model: facebookresearch/dinov2
+    version: dinov2_vitl14_reg
+    size: 518
+
+dense_scheduler:
+  target: diffusers.schedulers.FlowMatchEulerDiscreteScheduler
+  params:
+    num_train_timesteps: 1000
+    shift: 6.0 
+
+sparse_vae_512:
+  target: direct3d_s2.models.autoencoders.ss_vae.SparseSDFVAE
+  params:
+    use_checkpoint: true
+    embed_dim: 16
+    num_head_channels_encoder: 64
+    model_channels_encoder: 512 
+    num_heads_encoder: 8 
+    num_blocks_encoder: 4
+    num_head_channels_decoder: 64
+    model_channels_decoder: 512 
+    num_heads_decoder: 8 
+    num_blocks_decoder: 4
+    resolution: 64
+    out_channels: 1
+    use_fp16: true
+    latents_scale: 1.0
+    latents_shift: 0.0
+
+sparse_dit_512:
+  target: direct3d_s2.models.transformers.sparse_dit.SparseDiT 
+  params:
+    resolution: 64 
+    in_channels: 16
+    out_channels: 16
+    model_channels: 1024
+    cond_channels: 1024
+    num_blocks: 24
+    num_heads: 32
+    num_kv_heads: 2
+    compression_block_size: 8
+    selection_block_size: 8
+    topk: 8
+    compression_version: v1
+    pe_mode: ape
+    factor: 1.0
+    sparse_conditions: false
+    qk_rms_norm: true
+    use_shift: true
+    use_checkpoint: true
+    use_fp16: true
+
+sparse_scheduler_512:
+  target: diffusers.schedulers.FlowMatchEulerDiscreteScheduler
+  params:
+    num_train_timesteps: 1000
+    shift: 6.0
+
+sparse_vae_1024:
+  target: direct3d_s2.models.autoencoders.ss_vae.SparseSDFVAE
+  params:
+    use_checkpoint: true
+    embed_dim: 16
+    num_head_channels_encoder: 64
+    model_channels_encoder: 512 
+    num_heads_encoder: 8 
+    num_blocks_encoder: 4
+    num_head_channels_decoder: 64
+    model_channels_decoder: 512 
+    num_heads_decoder: 8 
+    num_blocks_decoder: 4
+    resolution: 128
+    out_channels: 1
+    use_fp16: true
+    latents_scale: 1.0
+    latents_shift: 0.0
+    chunk_size: 4
+
+sparse_dit_1024:
+  target: direct3d_s2.models.transformers.sparse_dit.SparseDiT 
+  params:
+    resolution: 128 
+    in_channels: 16
+    out_channels: 16
+    model_channels: 1024
+    cond_channels: 1024
+    num_blocks: 24
+    num_heads: 32
+    num_kv_heads: 2
+    compression_block_size: 8
+    selection_block_size: 8
+    topk: 8
+    compression_version: v1
+    pe_mode: ape
+    factor: 0.5
+    sparse_conditions: false
+    qk_rms_norm: true
+    use_shift: true
+    use_checkpoint: true
+    use_fp16: true
+
+sparse_scheduler_1024:
+  target: diffusers.schedulers.FlowMatchEulerDiscreteScheduler
+  params:
+    num_train_timesteps: 1000
+    shift: 8.0
+
+sparse_image_encoder:
+  target: direct3d_s2.models.conditioner.DinoEncoder
+  params:
+    model: facebookresearch/dinov2
+    version: dinov2_vitl14_reg
+    size: 518
+
+refiner:
+  target: direct3d_s2.models.refiner.unet_refiner.Voxel_RefinerXL
+  params:
+    in_channels: 1
+    out_channels: 1
+    layers_per_block:  2
+    layers_mid_block:  2
+    patch_size: 192
+    use_fp16: true
+
+refiner_1024:
+  target: direct3d_s2.models.refiner.unet_refiner.Voxel_RefinerXL_sign
+  params:
+    in_channels: 1
+    out_channels: 1
+    layers_per_block:  2
+    layers_mid_block:  2
+    patch_size: 256
+    use_fp16: true

direct3d-s2-v-1-1/model_dense.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5e6598155f907ca15eb9e4fb96905f9df4d10ebc319f9567ea947db8c684d8fd
+size 1505767932

direct3d-s2-v-1-1/model_refiner.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2eb0f70d323a43d63ab37c91790c8a7c2b9ada1d7771ffe99ef7a647bdbc99e4
+size 269562302

direct3d-s2-v-1-1/model_refiner_1024.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e75e405067f5c3eed48b2e46ee752e1f6593bf9a0db421cb4d694a677958460a
+size 269498584

direct3d-s2-v-1-1/model_sparse_1024.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c3349a91b1ee0d5843bfca4077e4d526f8e0856bd8a03c9f097b3312f386642d
+size 1094546938

direct3d-s2-v-1-1/model_sparse_512.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c54cd5527d6b040a931e00f7f6edd9460180726bce21cde03c7fe5c4f646e151
+size 1094546234