Comfy-Org / Comfy-Org/ComfyUI

`--use-ck-attention` silently corrupts TRELLIS.2 shape generation

Open
#16,027 0 comments 0 reactions 0 assignees View on GitHub
Potential Bug
Dominant language
Python
Stars
133k
Forks
15.7k
Avg merge
1d 7h
Merged PRs (30d)
158

Description

### Custom Node Testing

- [x] I have tried disabling custom nodes and the issue persists (see [how to disable custom nodes](https://docs.comfy.org/troubleshooting/custom-node-issues#step-1%3A-test-with-all-custom-nodes-disabled) if you need help)

### Expected Behavior

# `--use-ck-attention` silently corrupts TRELLIS.2 shape generation (garbage mesh, ~15x too few faces)

### Expected Behavior

Running the bundled `3d_pixal3d_trellis2_image_to_model` template with the TRELLIS.2 branch selected
(`Boolean (Switch to Trellis2)` = **true**, `trellis_2_int8_convrot.safetensors`) should produce a clean,
dense mesh of the input object — with the stock template input `viking_wolf_rune_axe.png` a solid axe of
roughly **7.2M vertices / 14.5M faces** at `GetMeshInfo` after `VaeDecodeShapeTrellis`.

### Actual Behavior

When ComfyUI is started with **`--use-ck-attention`**, the same workflow — same seeds, same models, same
everything — silently produces a **degenerate, spindly "skeleton" mesh**: sparse floating shells and
fragments that no longer resemble the input at all. `GetMeshInfo` reports only **~321K vertices / ~477K faces**
(about 1/15 of the correct geometry). No error, no warning, no traceback — the run reports `success`.

The DINOv3/structure stage still looks roughly plausible (the 32³ `VaeDecodeStructureTrellis2` voxel is a
recognizable blade+handle blob), so the corruption happens in the **shape / upsample stages**
(`Trellis2ShapeStage` → KSampler → `Trellis2UpsampleStage` → KSampler → `VaeDecodeShapeTrellis`).

| `--use-ck-attention` | no CK attention |
|---|---|
| ![broken](trellis2_with_ck_attention_BROKEN.jpg) | ![ok](trellis2_without_ck_attention_OK.jpg) |

### Steps to Reproduce

1. Install the TRELLIS.2 models for the template (`trellis_2_int8_convrot.safetensors`,
`trellis_2_shape_vae_bf16.safetensors`, `trellis_2_texture_vae_bf16.safetensors`,
`dino_v3_L_naf_fp32.safetensors`, `birefnet.safetensors`).
2. Load the template `3d_pixal3d_trellis2_image_to_model`, set `Boolean (Switch to Trellis2)` = `true`,
leave everything else at template defaults, input `viking_wolf_rune_axe.png`.
3. Start ComfyUI **without** `--use-ck-attention` → run. `GetMeshInfo`: **7,219,147 verts / 14,596,016 faces**,
output GLB is a correct textured axe.
4. Restart ComfyUI **with** `--use-ck-attention` → run the identical prompt.
`GetMeshInfo`: **321,072 verts / 476,806 faces**, output GLB is garbage.

Isolated by bisecting the server flags on one machine, one workflow, identical seeds
(`--disable-all-custom-nodes` in every test):

| server flags | GetMeshInfo verts / faces | result |
|---|---|---|
| *(none)* | 7,222,177 / 14,547,082 | ✅ correct |
| `--fast fp16_accumulation` | 7,219,147 / 14,596,016 | ✅ correct |
| `--use-ck-attention` | 507,367 / 787,938 | ❌ garbage |
| `--fast fp16_accumulation --use-ck-attention` | 321,072 / 476,806 | ❌ garbage |

(The exact numbers on the broken runs vary run to run; the shape is always garbage.)

An API-format copy of the exact prompt used for all four runs is attached as `repro_workflow_api.json`.

### Notes / possible cause

`--use-ck-attention` replaces the global `optimized_attention` with `attention_comfy_kitchen_int8`
(`comfy/ldm/modules/attention.py`, guarded by `model_management.comfy_kitchen_attention_enabled()`).
TRELLIS.2's sparse-structure transformer (`comfy/ldm/trellis2/`) evidently cannot tolerate INT8 attention:
the low-res structure stage survives, but the shape/upsample stages collapse into noise.

The practical problem is that this is a **global** switch with no per-model guard and it fails **silently** —
the run reports `execution_success` and produces a plausible-looking file, so a user with this flag in their
launch args (it is a big speedup on image DiTs) just concludes "TRELLIS.2 gives bad meshes".

It would help if TRELLIS.2 / Pixal3D either forced the default attention for their own blocks, or logged a
warning when INT8 attention is active for a model known to be incompatible with it.

### System Information

- ComfyUI commit: `3216c62e9962c3babd28a4dfea6e5aef50b8fe16` (2026-08-31), version 0.34.0
- OS: Ubuntu Linux, kernel 6.8
- Python 3.10.15, PyTorch 2.10.0+cu130
- GPU: NVIDIA GeForce RTX 3090 (sm_86), driver 595.84
- comfy-kitchen 0.2.31, comfyui-frontend-package 1.51.9, comfyui-workflow-templates 0.11.50
- Custom nodes disabled (`--disable-all-custom-nodes`) for every test above

[repro_workflow_api.json](https://github.com/user-attachments/files/31714292/repro_workflow_api.json)
Image
Image

### Actual Behavior

broken mesh

### Steps to Reproduce

Image

### Debug Logs

```powershell
No errors or warnings during execution; the run reports `execution_success`. The only relevant log line is at
startup:
```

### Other

_No response_

Contributor guide

Open the contributing guide

Research direction

Reproduce the bundled TRELLIS.2 workflow with and without --use-ck-attention, comparing GetMeshInfo results and execution logs. Trace the global optimized_attention selection in comfy/ldm/modules/attention.py and the shape and upsample stages under comfy/ldm/trellis2/. Done means TRELLIS.2 no longer silently produces corrupted meshes when the flag is enabled, either by using compatible attention or reporting the incompatibility.

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.