README.md

metadata

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

Presenting the first INT4 quantized efficient version of CubePart (arXiv 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, 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, 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

💡 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 torch.compile disabled, so dequantization runs in eager mode.

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 and the cube_part package from Roblox/cube. 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).

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:

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):

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, which is released under the OpenRAIL license. This derivative carries the same license.