Comfy-Org / Comfy-Org/comfy-kitchen

na3d custom-op registration uses PEP-585 annotations, breaking import on torch < 2.7

Open Beginner friendly
#114 1 comment 4 reactions 0 assignees View on GitHub
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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.