Comfy-Org / Comfy-Org/comfy-kitchen
na3d custom-op registration uses PEP-585 annotations, breaking import on torch < 2.7
- Dominant language
- Python
- Stars
- 220
- Forks
- 91
- Avg merge
- 1d 7h
- Merged PRs (30d)
- 12
Description
## Summary
Since **v0.2.28**, `import comfy_kitchen` fails outright on torch 2.6:
```
ValueError: infer_schema(func): Parameter kernel_size has unsupported type list[int].
```
The error is raised at package-import time, so the entire package becomes unimportable — not just the `na3d` op. Any downstream project on torch 2.6 breaks when it picks up a newer comfy-kitchen. That includes ComfyUI itself: its `requirements.txt` pins `comfy-kitchen==0.2.31` but leaves `torch` unpinned, so whether it works depends entirely on which torch happens to be installed.
## Root cause
`comfy_kitchen/backends/eager/na.py` (added in #95) registers the custom op with PEP-585 builtin-generic annotations:
```python
@torch.library.custom_op("comfy_kitchen::na3d", mutates_args=())
def _op_na3d(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
kernel_size: list[int],
is_causal: list[bool],
scale: float | None,
) -> torch.Tensor:
```
`torch.library.infer_schema` only learned to accept builtin generics in **torch 2.7.0**: `derived_seq_types()` in `torch/_library/infer_schema.py` gained `GenericAlias(list, (typ,))` and `GenericAlias(collections.abc.Sequence, (typ,))` there (pytorch/pytorch#146594). On 2.6 it accepts only `typing.List[int]` / `typing.Sequence[int]`.
Worth noting `scale: float | None` is *not* part of the problem — `float | None` compares and hashes equal to `typing.Optional[float]`, so the dict lookup resolves fine even on 2.6. Only the two `list[...]` parameters fail.
## Repro
torch 2.6.0+cu124, Python 3.10.12, Linux:
```python
>>> import comfy_kitchen
ValueError: infer_schema(func): Parameter kernel_size has unsupported type list[int]
```
Import chain: `comfy_kitchen/__init__.py:3` → `backends/cuda/__init__.py:161` → `backends/eager/__init__.py:71` → `backends/eager/na.py:164`.
## Suggested fix
Two annotations, no behavior change, restores torch 2.6 compatibility at zero cost:
```diff
- kernel_size: list[int],
- is_causal: list[bool],
+ kernel_size: typing.List[int],
+ is_causal: typing.List[bool],
```
Alternatively, if torch >= 2.7 is intended as a real floor, it would help a lot to declare it — `pyproject.toml` currently has `dependencies = []` with no mention of torch, so pip installs the package happily and the incompatibility only surfaces at runtime, with an error message that doesn't point at the version mismatch. Happy to send a PR for either direction.
## Environment
- comfy-kitchen 0.2.28 – 0.2.31 affected; 0.2.26 unaffected (predates #95)
- torch 2.6.0+cu124
- Python 3.10.12, Linux
Contributor guide
Research direction
Start in comfy_kitchen/backends/eager/na.py at the _op_na3d registration, then reproduce the import failure with torch 2.6.0. Verify that the two custom-op parameter annotations are accepted and that importing comfy_kitchen succeeds on torch 2.6 without changing behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python, pytorch
- Domain
- backend
- Issue type
- Bug
- Difficulty
- 1/5
- Estimated time
- Under an hour
- Activity status
- Quiet
- Clarity
- Clearly specified
- Newbie friendliness
- 85/100