Pixal3D bf16 and multiview checkpoints crash with tensor shape mismatch in k_rms_norm (int8_convrot works fine)
- Dominant language
- Python
- Stars
- 133k
- Forks
- 15.7k
- Avg merge
- 1d 7h
- Merged PRs (30d)
- 158
Description
## Bug description
Both `pixal3d_bf16.safetensors` and `pixal3d_multiview_int8_convrot.safetensors` crash during structure/shape sampling with a tensor shape mismatch inside Trellis2's cross-attention `k_rms_norm`. The default `pixal3d_int8_convrot.safetensors` checkpoint runs correctly with the exact same workflow (only the `UNETLoader` model file was swapped).
## Steps to reproduce
1. Load the official "Pixal3D & TRELLIS.2: Image to Model" workflow template (`3d_pixal3d_trellis2_image_to_model`).
2. Confirm it runs successfully with `pixal3d_int8_convrot.safetensors` loaded in the Pixal3D `UNETLoader` node (id 319 in the template).
3. Swap that node's `unet_name` widget to `pixal3d_bf16.safetensors` (no other graph changes). Queue the workflow with the same single input image.
4. Observe the crash below.
5. Separately: revert to `pixal3d_int8_convrot.safetensors`, but swap it to `pixal3d_multiview_int8_convrot.safetensors` instead, and feed a 2-image batch (front + side view, both 447x1407, combined via a standard `ImageBatch` node inserted before `RemoveBackground`/`ImageCropToMask`) into the same `Pixal3DConditioning` node. Same crash, at the same line.
## Error (both cases, identical traceback location)
```
File "ComfyUI\comfy\ldm\trellis2\model.py", line 519, in forward
k = self.k_rms_norm(k)
File ...\torch\nn\modules\module.py", line 1794, in _call_impl
return forward_call(*args, **kwargs)
File "ComfyUI\comfy\ldm\trellis2\model.py", line 80, in forward
return F.rms_norm(x, (x.shape[-1],)) * comfy.ops.cast_to_input(self.gamma, x)
RuntimeError: The size of tensor a (64) must match the size of tensor b (128) at non-singleton dimension 3
```
For the bf16 case this fails almost immediately (~6s into the prompt). For the multiview case it also fails early, during structure sampling, before any mesh is produced.
## Expected behavior
`pixal3d_bf16.safetensors` and `pixal3d_multiview_int8_convrot.safetensors` should run through the same graph (aside from whatever additional wiring multiview actually needs -- happy to be told the multiview path requires a different node graph than a plain `ImageBatch`, since I couldn't find any documented multiview workflow yet).
## Environment
- ComfyUI version: 0.34.3
- comfy-kitchen: 0.2.31, comfy-aimdo: 0.4.15
- torch: 2.14.0+cu130
- GPU: NVIDIA GeForce RTX 4070, 12GB VRAM
- OS: Windows 11
- Models used: `pixal3d_int8_convrot.safetensors` (works), `pixal3d_bf16.safetensors` (fails), `pixal3d_multiview_int8_convrot.safetensors` (fails), `dino_v3_L_naf_fp32.safetensors`, `trellis_2_shape_vae_bf16.safetensors`, `trellis_2_texture_vae_bf16.safetensors`, `birefnet.safetensors`, `moge_2_vitl_normal_fp16.safetensors` -- all downloaded from the official `Comfy-Org/Pixal3D`, `Comfy-Org/TRELLIS.2`, `Comfy-Org/BiRefNet`, `Comfy-Org/MoGe` repos.
## Additional notes
- `pixal3d_int8_convrot.safetensors` with the same template, same input image, works end-to-end (produces a valid textured GLB) once `DecimateMesh`'s `target_face_count` is lowered from the template's default 700000 to something more modest (700k caused an unrelated CUDA OOM during `UnwrapMesh`'s UV packing step on a 12GB card -- separate issue, not filing that one since it's more of a default-tuned-for-higher-VRAM-cards issue, but flagging in case it's relevant).
- Happy to provide the full workflow JSON or additional logs if useful.
Contributor guide
Research direction
Start with comfy/ldm/trellis2/model.py, especially the k_rms_norm call near line 519 and the RMS normalization near line 80. Reproduce the official 3d_pixal3d_trellis2_image_to_model workflow first with the working checkpoint, then compare tensor dimensions for the bf16 and multiview checkpoints at the failing path. Done means the affected checkpoints complete the reported workflow, or their required graph and compatibility constraints are identified and documented.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100