Comfy-Org / Comfy-Org/ComfyUI

Pixal3D bf16 and multiview checkpoints crash with tensor shape mismatch in k_rms_norm (int8_convrot works fine)

Open
#16,056 1 comment 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.