Add .safetensors checkpoint; README: Cube3D-style header, define HQQ (#3)

23cdd80 4 days ago

12.7 kB

---
license: openrail
base_model: Roblox/cubepart
base_model_relation: quantized
pipeline_tag: text-to-3d  
tags:
  - quantization
  - int4
  - hqq
  - torchao
  - 3d-generation
  - part-decomposition
  - 3d
  - mesh-generation
  - diffusion
---
# 🚀 First INT4 Quantized Efficient-CubePart3D - 83% Smaller Part Decomposition Model

![](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)

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.
| | 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
Measured on one NVIDIA L40S (44 GiB usable), PyTorch 2.8.0+cu128, torchao 0.11.0, eager mode, 50 diffusion steps, guidance scale 7.5, over 10 meshes, one from each of 10 categories. Two independent runs; ranges below are run 1 to run 2.
| Metric | BF16 | INT4 |
|---|---|---|
| Peak VRAM during generation | 33.71 GiB (both runs) | 27.09 GiB (both runs) |
| Latency per generation, mean of 10 meshes | 34.7 to 34.8 s | 66.5 to 66.9 s |
| Setup time (pipeline load) | 149.4 to 161.8 s | 153.5 to 166.1 s |
| Quantization time | n/a | 4.1 to 4.3 s |
| Diffusion transformer on disk | 8.582 GB | 1.479 GB |
Notes:
- Latency covers `input_to_part_shape` (diffusion sampling and mesh extraction). Surface sampling and shape encoding are excluded.
- Peak VRAM is `torch.cuda.max_memory_allocated()` over the full generation call. It includes the Qwen3-VL text encoder, the shape VAE and the mesh extraction step.
- INT4 setup time is the BF16 pipeline load plus quantization, because the INT4 file is applied on top of an instantiated pipeline (see Usage). A direct INT4-only load path is not implemented.
- Not tested on GPUs smaller than the L40S. With a measured 27.09 GiB peak, a 24 GB card is not expected to fit this workload as configured.
- Units: file sizes are decimal GB, as the file page shows them (exact bytes: 1,478,846,193 for the INT4 checkpoint and 8,582,452,032 for the BF16 diffusion transformer). VRAM figures are GiB (2^30 bytes) from `torch.cuda.max_memory_allocated()`.

## Quantization details

- Method: HQQ through `torchao.quantization.int4_weight_only(group_size=128, use_hqq=True)`, W4A16.
- Quantized: `nn.Linear` layers inside `transformer_blocks` and `multi_transformer_blocks` of `diffusion_model` (about 2134M of the 2145.6M Linear parameters in `diffusion_model`, 99.5%).
- Not quantized: `img_in`, `txt_in`, `proj_out`, `pos_embed`, `time_text_embed`, every layer whose name contains `norm`, the shape VAE (`shape_model`) and the Qwen3-VL text encoder (`base_model.text_encoder`).
- Why HQQ and not plain round to nearest: torchao's default tinygemm packing path raised a dtype mismatch ("Expected zeros dtype float32, got bfloat16") in the versions we tried (0.10 to 0.13). HQQ avoids that path and needs no calibration data.

## Quality benchmark (chamfer distance)

250 prompts across 12 categories, 828 parts in total (2 to 4 parts per prompt, part schemas written per prompt for this benchmark). Prompts come from the 310 prompt master suite in `TrNi/efficient-cube3d`; the categories `abstract_mathematical`, `geometric_primitive` and `symmetry_topology` are excluded.

Metric: average of the forward and backward mean nearest neighbour distance between 5000 seeded surface samples per part, in the units of the input meshes. All sampling is seeded, so the distance between two identical configurations is exactly 0.

- CD(A,B): same BF16 Cube3D input mesh, BF16 CubePart versus INT4 CubePart.
- CD(C,D): same INT4 Cube3D input mesh, BF16 CubePart versus INT4 CubePart.

| | min | median | mean | max |
|---|---|---|---|---|
| CD(A,B) | 0.0011 | 0.0114 | 0.0193 | 0.8206 |
| CD(C,D) | 0.0014 | 0.0118 | 0.0239 | 1.6261 |

Per category (n is the number of parts):

| Category | n | CD(A,B) median | CD(A,B) mean | CD(A,B) max | CD(C,D) median | CD(C,D) mean | CD(C,D) max |
|---|---|---|---|---|---|---|---|
| animal_domestic | 77 | 0.0100 | 0.0307 | 0.8206 | 0.0093 | 0.0249 | 0.7246 |
| animal_wild | 78 | 0.0106 | 0.0158 | 0.1654 | 0.0106 | 0.0220 | 0.5996 |
| architecture | 61 | 0.0162 | 0.0173 | 0.0541 | 0.0193 | 0.0270 | 0.3630 |
| electronics | 60 | 0.0141 | 0.0165 | 0.0835 | 0.0135 | 0.0197 | 0.2594 |
| fine_detail | 58 | 0.0142 | 0.0161 | 0.0512 | 0.0161 | 0.0245 | 0.0999 |
| furniture | 62 | 0.0131 | 0.0356 | 0.4662 | 0.0125 | 0.0153 | 0.0617 |
| musical_instrument | 69 | 0.0096 | 0.0177 | 0.1520 | 0.0110 | 0.0697 | 1.6261 |
| nature_plant | 61 | 0.0138 | 0.0167 | 0.1059 | 0.0134 | 0.0143 | 0.0541 |
| original_visuals | 107 | 0.0118 | 0.0203 | 0.3210 | 0.0113 | 0.0226 | 0.6126 |
| tool_hardware | 56 | 0.0094 | 0.0168 | 0.2186 | 0.0086 | 0.0200 | 0.2537 |
| vehicle_air_water | 68 | 0.0093 | 0.0130 | 0.1371 | 0.0081 | 0.0124 | 0.1168 |
| vehicle_land | 71 | 0.0097 | 0.0132 | 0.1608 | 0.0104 | 0.0127 | 0.1108 |

Reading the numbers: the median is about 0.011 for both pairs, and the mean sits above it because of a thin tail. The largest errors are small or thin parts next to a dominant part (for example a peacock head at 0.8206 in CD(A,B), a double bass endpin at 1.6261 and sitar frets at 1.2668 in CD(C,D)). CD(C,D) is higher than CD(A,B) on the overall statistics but lower in some categories, so it is a mild aggregate trend and not a per mesh rule.

## Comparison to Cube3D INT4

The `efficient-cube3d` card reports a median chamfer distance of 67.7e-3 for Cube3D INT4 over 15 categories and 310 prompts. Recomputing Cube3D's own BF16 versus INT4 error with the seeded protocol above, on the 250 meshes of the 12 shared categories, gives a median of 45.5e-3. CubePart INT4 under the same protocol has a per mesh median of 13.0e-3 (median of each mesh's mean over its parts), about 3.5 times lower, and lower in every one of the 12 categories (ratio 2.2 to 5.3).

These measure different things: whole mesh drift for Cube3D and per part decomposition drift on a fixed input mesh for CubePart. Treat the ratio as indicative, not exact. The two models also fail on different meshes (per mesh rank correlation 0.27, and 2 of the 15 worst meshes in common).

## Usage

You need the base weights and code from [Roblox/cubepart](https://huggingface.co/Roblox/cubepart) and the `cube_part` package from [Roblox/cube](https://github.com/Roblox/cube/tree/main/cubepart). The INT4 file replaces the diffusion transformer weights on an instantiated pipeline. The order matters: quantize first, then load the state dict (no `assign=True`).

```python
import torch
from torchao.quantization import quantize_, int4_weight_only
from cube_part.pipelines import PartShapeDenoiserPipeline
pipe = PartShapeDenoiserPipeline(
    config_path="configs/shape_denoiser_multimesh.yaml",
    checkpoint_path="weights/multi_part_dit.safetensors",   # base BF16 weights
    vae_checkpoint_path="weights/vae.safetensors",
    extract_geometry_fn_name="extract_geometry_coarse_to_fine",
)
SKIP = ("img_in", "txt_in", "proj_out", "pos_embed", "time_text_embed", "norm")
def only_blocks(module, fqn):
    return isinstance(module, torch.nn.Linear) and not any(s in fqn for s in SKIP)
quantize_(pipe.system.diffusion_model,
          int4_weight_only(group_size=128, use_hqq=True),
          filter_fn=only_blocks)
sd = torch.load("multi_part_dit_int4_hqq_g128.pt", map_location="cuda", weights_only=False)
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
import numpy as np, torch, trimesh
from cube_part.pipelines import ShapeInput
from cube_part.utils.mesh import sample_surface
mesh = trimesh.load("input.glb", force="mesh")
parts = ["circular seat", "vertical legs", "footrest ring"]
np.random.seed(0)
surface = sample_surface(mesh, num_samples=128_000)
surface = torch.from_numpy(surface).to(pipe.device).unsqueeze(0).float()
latents, _ = pipe.encode_shape(surface)
part_meshes = pipe.input_to_part_shape(
    ShapeInput(prompt=[parts], latents=latents),
    guidance_scale=7.5, num_inference_steps=50, seed=0,
)
# part_meshes[i] is a (vertices, faces) pair for parts[i]
```

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 `.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: 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 (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.

## Limitations

- INT4 is slower than BF16 here (eager mode). We expect `torch.compile` accounts for much of the gap but have not tested compiled execution.
- Peak VRAM is 27.09 GiB; smaller GPUs were not tested.
- The base BF16 weights are needed to instantiate the pipeline. There is no INT4 only loading path yet.
- One quantization setting only (HQQ W4A16, group size 128). No other bit widths or methods were compared.
- The error of a fully INT4 pipeline (INT4 Cube3D input and INT4 CubePart against a BF16 reference) was not measured. CD(C,D) only measures what CubePart's quantization adds on an INT4 input.
- Chamfer distance is a point cloud metric. Watertightness and part reassembly were not checked.

## License

Derived from [Roblox/cubepart](https://huggingface.co/Roblox/cubepart), which is released under the OpenRAIL license. This derivative carries the same license.

	---
	license: openrail
	base_model: Roblox/cubepart
	base_model_relation: quantized
	pipeline_tag: text-to-3d
	tags:
	- quantization
	- int4
	- hqq
	- torchao
	- 3d-generation
	- part-decomposition
	- 3d
	- mesh-generation
	- diffusion
	---

	# 🚀 First INT4 Quantized Efficient-CubePart3D - 83% Smaller Part Decomposition Model

	![](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)

	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.

	\| \| 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

	Measured on one NVIDIA L40S (44 GiB usable), PyTorch 2.8.0+cu128, torchao 0.11.0, eager mode, 50 diffusion steps, guidance scale 7.5, over 10 meshes, one from each of 10 categories. Two independent runs; ranges below are run 1 to run 2.

	\| Metric \| BF16 \| INT4 \|
	\|---\|---\|---\|
	\| Peak VRAM during generation \| 33.71 GiB (both runs) \| 27.09 GiB (both runs) \|
	\| Latency per generation, mean of 10 meshes \| 34.7 to 34.8 s \| 66.5 to 66.9 s \|
	\| Setup time (pipeline load) \| 149.4 to 161.8 s \| 153.5 to 166.1 s \|
	\| Quantization time \| n/a \| 4.1 to 4.3 s \|
	\| Diffusion transformer on disk \| 8.582 GB \| 1.479 GB \|

	Notes:

	- Latency covers `input_to_part_shape` (diffusion sampling and mesh extraction). Surface sampling and shape encoding are excluded.
	- Peak VRAM is `torch.cuda.max_memory_allocated()` over the full generation call. It includes the Qwen3-VL text encoder, the shape VAE and the mesh extraction step.
	- INT4 setup time is the BF16 pipeline load plus quantization, because the INT4 file is applied on top of an instantiated pipeline (see Usage). A direct INT4-only load path is not implemented.
	- Not tested on GPUs smaller than the L40S. With a measured 27.09 GiB peak, a 24 GB card is not expected to fit this workload as configured.
	- Units: file sizes are decimal GB, as the file page shows them (exact bytes: 1,478,846,193 for the INT4 checkpoint and 8,582,452,032 for the BF16 diffusion transformer). VRAM figures are GiB (2^30 bytes) from `torch.cuda.max_memory_allocated()`.

	## Quantization details

	- Method: HQQ through `torchao.quantization.int4_weight_only(group_size=128, use_hqq=True)`, W4A16.
	- Quantized: `nn.Linear` layers inside `transformer_blocks` and `multi_transformer_blocks` of `diffusion_model` (about 2134M of the 2145.6M Linear parameters in `diffusion_model`, 99.5%).
	- Not quantized: `img_in`, `txt_in`, `proj_out`, `pos_embed`, `time_text_embed`, every layer whose name contains `norm`, the shape VAE (`shape_model`) and the Qwen3-VL text encoder (`base_model.text_encoder`).
	- Why HQQ and not plain round to nearest: torchao's default tinygemm packing path raised a dtype mismatch ("Expected zeros dtype float32, got bfloat16") in the versions we tried (0.10 to 0.13). HQQ avoids that path and needs no calibration data.

	## Quality benchmark (chamfer distance)

	250 prompts across 12 categories, 828 parts in total (2 to 4 parts per prompt, part schemas written per prompt for this benchmark). Prompts come from the 310 prompt master suite in `TrNi/efficient-cube3d`; the categories `abstract_mathematical`, `geometric_primitive` and `symmetry_topology` are excluded.

	Metric: average of the forward and backward mean nearest neighbour distance between 5000 seeded surface samples per part, in the units of the input meshes. All sampling is seeded, so the distance between two identical configurations is exactly 0.

	- CD(A,B): same BF16 Cube3D input mesh, BF16 CubePart versus INT4 CubePart.
	- CD(C,D): same INT4 Cube3D input mesh, BF16 CubePart versus INT4 CubePart.

	\| \| min \| median \| mean \| max \|
	\|---\|---\|---\|---\|---\|
	\| CD(A,B) \| 0.0011 \| 0.0114 \| 0.0193 \| 0.8206 \|
	\| CD(C,D) \| 0.0014 \| 0.0118 \| 0.0239 \| 1.6261 \|

	Per category (n is the number of parts):

	\| Category \| n \| CD(A,B) median \| CD(A,B) mean \| CD(A,B) max \| CD(C,D) median \| CD(C,D) mean \| CD(C,D) max \|
	\|---\|---\|---\|---\|---\|---\|---\|---\|
	\| animal_domestic \| 77 \| 0.0100 \| 0.0307 \| 0.8206 \| 0.0093 \| 0.0249 \| 0.7246 \|
	\| animal_wild \| 78 \| 0.0106 \| 0.0158 \| 0.1654 \| 0.0106 \| 0.0220 \| 0.5996 \|
	\| architecture \| 61 \| 0.0162 \| 0.0173 \| 0.0541 \| 0.0193 \| 0.0270 \| 0.3630 \|
	\| electronics \| 60 \| 0.0141 \| 0.0165 \| 0.0835 \| 0.0135 \| 0.0197 \| 0.2594 \|
	\| fine_detail \| 58 \| 0.0142 \| 0.0161 \| 0.0512 \| 0.0161 \| 0.0245 \| 0.0999 \|
	\| furniture \| 62 \| 0.0131 \| 0.0356 \| 0.4662 \| 0.0125 \| 0.0153 \| 0.0617 \|
	\| musical_instrument \| 69 \| 0.0096 \| 0.0177 \| 0.1520 \| 0.0110 \| 0.0697 \| 1.6261 \|
	\| nature_plant \| 61 \| 0.0138 \| 0.0167 \| 0.1059 \| 0.0134 \| 0.0143 \| 0.0541 \|
	\| original_visuals \| 107 \| 0.0118 \| 0.0203 \| 0.3210 \| 0.0113 \| 0.0226 \| 0.6126 \|
	\| tool_hardware \| 56 \| 0.0094 \| 0.0168 \| 0.2186 \| 0.0086 \| 0.0200 \| 0.2537 \|
	\| vehicle_air_water \| 68 \| 0.0093 \| 0.0130 \| 0.1371 \| 0.0081 \| 0.0124 \| 0.1168 \|
	\| vehicle_land \| 71 \| 0.0097 \| 0.0132 \| 0.1608 \| 0.0104 \| 0.0127 \| 0.1108 \|

	Reading the numbers: the median is about 0.011 for both pairs, and the mean sits above it because of a thin tail. The largest errors are small or thin parts next to a dominant part (for example a peacock head at 0.8206 in CD(A,B), a double bass endpin at 1.6261 and sitar frets at 1.2668 in CD(C,D)). CD(C,D) is higher than CD(A,B) on the overall statistics but lower in some categories, so it is a mild aggregate trend and not a per mesh rule.

	## Comparison to Cube3D INT4

	The `efficient-cube3d` card reports a median chamfer distance of 67.7e-3 for Cube3D INT4 over 15 categories and 310 prompts. Recomputing Cube3D's own BF16 versus INT4 error with the seeded protocol above, on the 250 meshes of the 12 shared categories, gives a median of 45.5e-3. CubePart INT4 under the same protocol has a per mesh median of 13.0e-3 (median of each mesh's mean over its parts), about 3.5 times lower, and lower in every one of the 12 categories (ratio 2.2 to 5.3).

	These measure different things: whole mesh drift for Cube3D and per part decomposition drift on a fixed input mesh for CubePart. Treat the ratio as indicative, not exact. The two models also fail on different meshes (per mesh rank correlation 0.27, and 2 of the 15 worst meshes in common).

	## Usage

	You need the base weights and code from [Roblox/cubepart](https://huggingface.co/Roblox/cubepart) and the `cube_part` package from [Roblox/cube](https://github.com/Roblox/cube/tree/main/cubepart). The INT4 file replaces the diffusion transformer weights on an instantiated pipeline. The order matters: quantize first, then load the state dict (no `assign=True`).

	```python
	import torch
	from torchao.quantization import quantize_, int4_weight_only
	from cube_part.pipelines import PartShapeDenoiserPipeline

	pipe = PartShapeDenoiserPipeline(
	config_path="configs/shape_denoiser_multimesh.yaml",
	checkpoint_path="weights/multi_part_dit.safetensors", # base BF16 weights
	vae_checkpoint_path="weights/vae.safetensors",
	extract_geometry_fn_name="extract_geometry_coarse_to_fine",
	)

	SKIP = ("img_in", "txt_in", "proj_out", "pos_embed", "time_text_embed", "norm")
	def only_blocks(module, fqn):
	return isinstance(module, torch.nn.Linear) and not any(s in fqn for s in SKIP)

	quantize_(pipe.system.diffusion_model,
	int4_weight_only(group_size=128, use_hqq=True),
	filter_fn=only_blocks)

	sd = torch.load("multi_part_dit_int4_hqq_g128.pt", map_location="cuda", weights_only=False)
	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
	import numpy as np, torch, trimesh
	from cube_part.pipelines import ShapeInput
	from cube_part.utils.mesh import sample_surface

	mesh = trimesh.load("input.glb", force="mesh")
	parts = ["circular seat", "vertical legs", "footrest ring"]

	np.random.seed(0)
	surface = sample_surface(mesh, num_samples=128_000)
	surface = torch.from_numpy(surface).to(pipe.device).unsqueeze(0).float()
	latents, _ = pipe.encode_shape(surface)

	part_meshes = pipe.input_to_part_shape(
	ShapeInput(prompt=[parts], latents=latents),
	guidance_scale=7.5, num_inference_steps=50, seed=0,
	)
	# part_meshes[i] is a (vertices, faces) pair for parts[i]
	```

	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 `.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: 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 (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.

	## Limitations

	- INT4 is slower than BF16 here (eager mode). We expect `torch.compile` accounts for much of the gap but have not tested compiled execution.
	- Peak VRAM is 27.09 GiB; smaller GPUs were not tested.
	- The base BF16 weights are needed to instantiate the pipeline. There is no INT4 only loading path yet.
	- One quantization setting only (HQQ W4A16, group size 128). No other bit widths or methods were compared.
	- The error of a fully INT4 pipeline (INT4 Cube3D input and INT4 CubePart against a BF16 reference) was not measured. CD(C,D) only measures what CubePart's quantization adds on an INT4 input.
	- Chamfer distance is a point cloud metric. Watertightness and part reassembly were not checked.

	## License

	Derived from [Roblox/cubepart](https://huggingface.co/Roblox/cubepart), which is released under the OpenRAIL license. This derivative carries the same license.