Add .safetensors checkpoint; README: Cube3D-style header, define HQQ

README.md

@@ -15,16 +15,28 @@ tags:
   - diffusion
 ---
 
-# CubePart INT4 (HQQ W4A16, group size 128)
+# 🚀 First INT4 Quantized Efficient-CubePart3D - 83% Smaller Part Decomposition Model
 
-INT4 weight-only quantization of the diffusion transformer in Roblox's [CubePart](https://huggingface.co/Roblox/cubepart) (open-vocabulary, part-controllable 3D mesh generation, [arXiv 2605.28763](https://arxiv.org/abs/2605.28763), SIGGRAPH 2026). Companion to [TrNi/efficient-cube3d](https://huggingface.co/TrNi/efficient-cube3d), which does the same for Cube3D.
+![](https://img.shields.io/badge/Model_Size-1.48_GB_(82.8%25↓)-blueviolet)
+![](https://img.shields.io/badge/Peak_VRAM-27.09_GiB_(19.6%25↓)-brightgreen)
+![](https://img.shields.io/badge/Part_CD_median-0.0114_(near_lossless)-blue)
+![](https://img.shields.io/badge/Latency-66.5s_(eager_mode,_1.9x_slower)-orange)
 
-## Summary
+Presenting the **first INT4 quantized efficient version** of [CubePart](https://huggingface.co/Roblox/cubepart) ([arXiv 2605.28763](https://arxiv.org/abs/2605.28763), SIGGRAPH 2026), an open-vocabulary, part-controllable 3D mesh generative model.
+Quantized via **HQQ (Half-Quadratic Quantization)**, a calibration-free weight-only method, at **W4A16** (4-bit weights, 16-bit activations, group_size=128) using [torchao](https://github.com/pytorch/ao),
+it cuts the diffusion transformer from **8.58 GB → 1.48 GB (82.8%↓)** and peak VRAM from **33.71 GiB → 27.09 GiB (19.6%↓)** while keeping part decomposition near lossless (median chamfer distance **0.0114** per part over 828 parts).
+Companion to [TrNi/efficient-cube3d](https://huggingface.co/TrNi/efficient-cube3d), which does the same for Cube3D.
 
-- Checkpoint size: 1.479 GB, versus 8.582 GB for the BF16 diffusion transformer (82.8% smaller).
-- Peak VRAM during generation: 27.09 GiB, versus 33.71 GiB for BF16 (19.6% lower).
-- Quality: median chamfer distance between BF16 and INT4 part outputs is 0.0114 per part, which is near lossless on the typical part. Errors concentrate in a few small, thin parts.
-- Latency: INT4 is slower than BF16 in this release (66.5 to 66.9 s versus 34.7 to 34.8 s per generation). `torch.compile` has to be disabled for CubePart, so INT4 dequantization runs in eager mode.
+| | BF16 | **INT4 (HQQ W4A16)** |
+|---|:-:|:-:|
+| 💾 Model size (diffusion transformer) | 8.58 GB | **1.48 GB (82.8%↓)** |
+| 🎮 Peak VRAM | 33.71 GiB | **27.09 GiB (19.6%↓)** ✨ |
+| 📦 Setup time | 161.8 s | **166.1 s** (BF16 load + 4.3 s quantize) |
+| ⏱️ Latency | 34.7 s | **66.5 s** (eager mode) |
+| 🎯 Part chamfer distance vs BF16, median | n/a | **0.0114** |
+
+<mark>💡 The diffusion transformer is 82.8% smaller on disk and part decomposition stays near lossless on the typical part. Peak VRAM drops less (19.6%) because the Qwen3-VL text encoder and the shape VAE stay in BF16. INT4 is slower here because CubePart needs <code>torch.compile</code> disabled, so dequantization runs in eager mode.
+</mark>
 
 ## Efficiency
 
@@ -120,6 +132,29 @@ sd = torch.load("multi_part_dit_int4_hqq_g128.pt", map_location="cuda", weights_
 pipe.system.diffusion_model.load_state_dict(sd)
 ```
 
+Or load the `.safetensors` file instead (no pickle). safetensors stores plain tensors only, so each INT4 tensor was saved as its inner tensors under keys like `<param>::tensor_impl::packed_weight`. After the same `quantize_` call above, copy them back in:
+
+```python
+from safetensors.torch import load_file
+
+def leaves(prefix, t, out):
+    # walk torchao tensor subclasses down to their plain inner tensors
+    if hasattr(t, "__tensor_flatten__"):
+        for name in t.__tensor_flatten__()[0]:
+            leaves(f"{prefix}::{name}", getattr(t, name), out)
+    else:
+        out[prefix] = t
+
+saved = load_file("multi_part_dit_int4_hqq_g128.safetensors", device="cuda")
+flat = {}
+for k, v in pipe.system.diffusion_model.state_dict(keep_vars=True).items():
+    leaves(k, v, flat)
+assert set(flat) == set(saved), "keys do not match: check the quantize_ settings"
+with torch.no_grad():
+    for k, t in flat.items():
+        t.copy_(saved[k])
+```
+
 Generating parts from a mesh (this is the call used for every number above):
 
 ```python
@@ -144,15 +179,16 @@ part_meshes = pipe.input_to_part_shape(
 
 Environment notes: set `TORCHDYNAMO_DISABLE=1` (CubePart's compiled path raises a Dynamo error in `hijack_qwenimage_multi.py`). On our setup `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` also avoided a Warp marching cubes out of memory error.
 
-The checkpoint is a pickle and needs `weights_only=False` because torchao's INT4 tensors are tensor subclasses. Only load it from a source you trust.
+The `.pt` checkpoint is a pickle and needs `weights_only=False` because torchao's INT4 tensors are tensor subclasses. Only load it from a source you trust, or use the `.safetensors` file, which contains no code.
 
-Reload check: a fresh process that quantizes and then loads this file reproduces a live quantized model's output exactly (vertices and faces bit for bit on all 3 parts of a bar stool prompt).
+Reload check: for each file, a fresh process that quantizes and then loads it reproduces a live quantized model's output exactly (vertices and faces bit for bit on all 3 parts of a bar stool prompt). Both files hold the same weights.
 
 ## Files
 
 | File | Size | Description |
 |---|---|---|
-| `multi_part_dit_int4_hqq_g128.pt` | 1.479 GB | INT4 `diffusion_model` state dict, saved after in place quantization |
+| `multi_part_dit_int4_hqq_g128.pt` | 1.479 GB | INT4 `diffusion_model` state dict, saved after in place quantization (pickle) |
+| `multi_part_dit_int4_hqq_g128.safetensors` | 1.478 GB | Same weights as plain tensors, 1163 entries (no pickle) |
 | `quant_config.json` | small | Quantization settings and load recipe |
 
 The shape VAE and the text encoder are unchanged, so they are not re-hosted. Use the base repository for them.

multi_part_dit_int4_hqq_g128.safetensors

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

quant_config.json

@@ -8,11 +8,23 @@
   "config": "int4_weight_only(group_size=128, use_hqq=True)",
   "group_size": 128,
   "quantized_scope": "nn.Linear layers in transformer_blocks and multi_transformer_blocks of diffusion_model",
-  "skipped_name_substrings": ["img_in", "txt_in", "proj_out", "pos_embed", "time_text_embed", "norm"],
-  "untouched_modules": ["shape_model", "base_model.text_encoder"],
+  "skipped_name_substrings": [
+    "img_in",
+    "txt_in",
+    "proj_out",
+    "pos_embed",
+    "time_text_embed",
+    "norm"
+  ],
+  "untouched_modules": [
+    "shape_model",
+    "base_model.text_encoder"
+  ],
   "checkpoint_file": "multi_part_dit_int4_hqq_g128.pt",
+  "checkpoint_file_safetensors": "multi_part_dit_int4_hqq_g128.safetensors",
   "checkpoint_contents": "diffusion_model.state_dict() after in place quantization",
   "load_recipe": "instantiate the pipeline, quantize_ diffusion_model with the same config and filter, then load_state_dict with defaults (no assign=True), weights_only=False",
+  "load_recipe_safetensors": "instantiate the pipeline, quantize_ diffusion_model with the same config and filter, flatten each state_dict tensor via __tensor_flatten__ into '<key>::<inner>' leaves, then copy_ each leaf from the file (see README)",
   "required_env": {
     "TORCHDYNAMO_DISABLE": "1"
   }