feat(deps): resolve torch from PyPI on linux/aarch64 to enable ARM64 installs and Docker builds

[stellarfeline]

Summary

Makes InvokeAI installable (and the standard docker/Dockerfile buildable) on Linux ARM64 / aarch64, using per-dependency platform markers in pyproject.toml — no separate Dockerfile and no lockfile bypass.

The problem: the pinned +cpu / +cu128 / +rocm7.1 torch+torchvision pairs live on PyTorch's WHL indexes. Those indexes carry aarch64 wheels for torch (+cpu, +cu128), but no aarch64 torchvision wheels, so the pinned pairs are unsatisfiable on linux/aarch64 and installation fails for every extra.

The fix:

• pyproject.toml
  • The [tool.uv.sources] entries for torch / torchvision gain a marker excluding linux/aarch64 from the WHL indexes, plus an explicit PyPI source entry for that platform.
  • The +cpu / +cu128 / +rocm7.1 pins under the cpu / cuda / rocm extras are gated behind the same marker, with explicit PyPI-resolved fallback pins (torch==2.7.1, torchvision==0.22.1) for linux/aarch64. The explicit fallback entries keep each extra's conflict-universe populated on aarch64 and make the resolution robust across uv versions.
  • pytorch-triton-rocm→triton-rocm source marker tightened to exclude aarch64.
  • [tool.uv] environments extended to include linux/aarch64.
• uv.lock regenerated; non-aarch64 resolution paths are unchanged.
• docker/README.md / docker/Dockerfile: the "only x86_64 is supported" note is corrected, with arm64 build instructions (including docker buildx --platform linux/arm64 cross-compilation) and a refreshed comment.
• Site docs (docs/src/content/docs/): ARM64 CPU-only support noted in the hardware requirements; the manual install guide warns against --torch-backend=cpu on aarch64 (that index has no aarch64 torchvision wheels); the dev environment guide notes --extra cpu for ARM64; the Docker configuration page documents GPU_DRIVER=cpu.

Behavior note: on linux/aarch64, every extra (including cuda and rocm) installs the CPU-only PyPI wheels instead of failing loudly. Since torch+cu128 aarch64 wheels do exist on the index, real CUDA-on-aarch64 (Grace/GH200) is feasible later with a torchvision-specific workaround — breadcrumbs are in the pyproject.toml comments.

Verified

• Native: Linux x86_64 (cpu, cuda), Windows AMD64 (cpu, cuda), Linux aarch64 (cpu), macOS arm64.
• Docker: Linux aarch64 (cpu) via the standard docker/Dockerfile with --build-arg GPU_DRIVER=cpu (Raspberry Pi 5, Ubuntu 24.04).
• uv lock --locked passes; uv sync --frozen --extra cpu + torch/torchvision import verified on aarch64 after the merge with main.

QA Instructions

On any linux/aarch64 machine:

uv sync --frozen --extra cpu
# or, with Docker, from the repo root:
docker build -f docker/Dockerfile --build-arg GPU_DRIVER=cpu -t invokeai:cpu .
docker run --publish 9090:9090 --volume ~/invokeai:/invokeai invokeai:cpu

Navigate to http://localhost:9090 and verify the UI loads. API-backed models (GPT Image, Gemini) can be configured and used without a GPU.

On x86_64 / Windows / macOS, resolution is unchanged — uv sync --frozen --extra <cpu|cuda|rocm> behaves exactly as on main.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable)
• Documentation added / updated (if applicable)

[lstein]

● Code Review Summary — feat/docker-arm64

The approach is sound and the awk/sed patching logic correct for the current pyproject.toml. The layer caching is well-structured, and the documentation updates are greatly appreciated.

No security issues.

A few changes are requested:

Important:

• Patching is fragile — silently produces incorrect output if pyproject.toml format changes. Add a pre-flight guard like grep -q '^torch = [' before the awk pipeline.
• No CI coverage: build-container.yml only triggers on docker/Dockerfile and builds linux/amd64 only. Consider adding a path trigger for Dockerfile.arm64.

Suggestions:

• torchvision ends up unpinned ("torchvision"); consider ~=0.22.0 to match the torch~=2.7.0 constraint.
• The modified pyproject.toml and uv.lock stay baked into the final image (because COPY is used instead of the main Dockerfile's --mount=type=bind). Add RUN rm pyproject.toml
  uv.lock post-install.
• README could mention docker buildx --platform linux/arm64 for cross-compilation from x86_64.

[keturn]

It's not ideal to have a target that only builds with a specific Dockerfile. If pyproject.toml doesn't work, how would someone develop on this this platform? How could it be published to PyPI or some other non-Docker registry? How to avoid unplanned dependency changes if it doesn't use the lockfile?

It'd be preferable to make sure platform-specific dependencies have appropriate platform markers in pyproject.toml so they're appropriately resolved and written to the universal lockfile.

This is definitely non-trivial dependency management, but hopefully uv is up to the task:

• https://docs.astral.sh/uv/concepts/resolution/#platform-markers
• https://docs.astral.sh/uv/guides/integration/pytorch/#configuring-accelerators-with-optional-dependencies

[stellarfeline]

Looking into that

[stellarfeline]

Rewrote the commit with @keturn 's approach, now it builds both locally and within docker on Linux ARM64 platforms.

• Local build uv sync --frozen --extra cpu
• Docker build docker build -f docker/Dockerfile --build-arg GPU_DRIVER=cpu -t invokeai:arm64-cpu .

Re: @lstein's earlier review — since the dedicated Dockerfile is gone, the patching-related concerns and the CI path-trigger suggestion no longer apply. The README cross-compile note is included.

[keturn]

That looks cleaner!

As far as I can tell, mac uses the default index as well, so that should continue to work.

We should find someone to test on Windows. According to uv docs platform_machine might be different there.

[stellarfeline]

I have a Windows PC, I can test it later.

[stellarfeline]

Force-pushed to replace the previous two commits with a single commit. Summary of what changed:

pyproject.toml

The three PyTorch WHL-index [tool.uv.sources] entries previously used platform_machine == 'x86_64' (then extended to 'x86_64' or 'AMD64' in the follow-up commit). They now use:

sys_platform != 'linux' or platform_machine != 'aarch64'

The same marker is applied to the +cpu / +cu128 / +rocm6.3 version pins under [project.optional-dependencies] (cpu, cuda, rocm).

pytorch-triton-rocm similarly gains and platform_machine != 'aarch64' alongside its existing sys_platform == 'linux' constraint.

The marker resolves to True on Linux x86_64, Windows (regardless of AMD64 vs. x86_64), and macOS — so the custom indexes and pinned +suffix wheels apply unchanged on those platforms. It resolves to False only on Linux aarch64, where uv then falls back to the base torch~=2.7.0 / torchvision
declarations from PyPI.

uv.lock

Regenerated to match. Non-aarch64 resolution paths are unchanged from main.

Verified builds

• Native: Linux x86_64 (cpu, cuda), Windows AMD64 (cpu, cuda), Linux aarch64 (cpu), macOS arm64.
• Docker: Linux aarch64 (cpu) via the standard docker/Dockerfile with --build-arg GPU_DRIVER=cpu.

[lstein]

First off — my sincere apologies for how long this review took. You submitted this a couple of months ago and it deserved attention much sooner. Thank you for your patience, and for the work you put into this.

The good news: the current marker-based approach is a big improvement over the original Dockerfile.arm64 revision. It makes the existing docker/Dockerfile work on arm64 (via --build-arg GPU_DRIVER=cpu), avoids a second Dockerfile to keep in sync, and correctly uses per-dependency markers instead of a global tool.uv.environments filter. I verified that uv lock --locked passes on the branch head, so the CI lockfile gate is satisfied, and the lock diff is consistent with a clean uv lock regeneration. The marker logic itself checks out — it's a no-op on macOS (platform_machine == 'arm64') and Windows-on-ARM ('ARM64'), and the tightened pytorch-triton-rocm source marker is the correct complement.

I'm requesting changes only for the items below, which are mostly metadata/docs rather than the mechanism itself:

1. PR title and description are stale

The title says "feat(docker): add Dockerfile.arm64" and the body describes the earlier awk/sed-patched Dockerfile approach — but the actual diff contains no docker changes at all; it's pyproject.toml markers + a regenerated uv.lock. Please retitle (e.g. feat(deps): resolve torch from PyPI on linux/aarch64) and rewrite the body to describe the marker approach, since the title feeds our changelog.

2. Docs still say arm64 is unsupported

• docker/README.md line 91 still reads "only x86_64 architecture is supported." Since this PR is what makes arm64 builds possible, please update it here and add the one-liner build instruction (docker build --build-arg GPU_DRIVER=cpu ...).
• The comment at docker/Dockerfile line 69 (# NOTE: there are no pytorch builds for arm64 + cuda, only cpu) is now only half-true and worth refreshing (see item 3).

3. The rationale comment in pyproject.toml is factually off

The new comment says the PyTorch WHL indexes "only publish linux_x86_64 / win_amd64." In fact the lockfile itself contains torch-2.7.1+cpu and torch-2.7.1+cu128 manylinux_2_28_aarch64 wheels — the real blocker is torchvision (+cpu / +cu128 / +rocm6.3 have no aarch64 wheels). Please correct the comment so a future maintainer doesn't "fix" the markers based on a wrong premise.

4. Document the silent CPU fallback for --extra cuda on aarch64

On Grace/GH200-class arm64 hosts, uv sync --extra cuda previously failed loudly; with this change it silently installs the CPU-only PyPI torch wheel. That's an acceptable trade (it never worked before), but please note it in the docs/comment. Since torch+cu128 aarch64 wheels do exist, real CUDA-on-arm support is feasible later with a torchvision-only workaround — a comment breadcrumb would help whoever picks that up.

Optional (non-blocking)

Nothing in CI exercises aarch64 resolution, so a future dependency bump could silently re-break this. A cheap guard that needs no ARM hardware would be a CI step running uv pip compile pyproject.toml --extra cpu --python-platform aarch64-unknown-linux-gnu (resolution-only). Happy to take that as a follow-up PR instead.

Thanks again for the contribution and for reworking it into the cleaner approach — once the description and docs are aligned with the actual diff, this is good to go.

[lstein]

Please see comment above.

[lstein]

There were merge conflicts in pyproject.toml, so I took the liberty of resolving them and pushing the fix to your branch. Hope you don't mind.

Details:

• Merged current main (ROCm 7.1 + Transformers 5.5.4) into the branch and regenerated uv.lock rather than hand-merging it.
• Main's new tool.uv.environments setting excluded linux/aarch64 from resolution entirely — it's been extended to include it, otherwise the PR would have silently become a no-op.
• Two uv subtleties required additions beyond her original markers: explicit PyPI-pinned torch==2.7.1/torchvision==0.22.1 fallback entries inside each of the cpu/cuda/rocm extras (uv's conflict-extras otherwise leave aarch64 with no torch at all), and an explicit pypi index source for torch/torchvision on aarch64 (uv doesn't fall back to the default index inside an extra's resolution split).
• Verified from the lock: all extras on linux/aarch64 resolve PyPI torch 2.7.1 + torchvision 0.22.1; x86_64/macOS/Windows are unchanged from main. Worth re-testing on your Pi 5 to confirm end-to-end.
• The remedial items from my review (PR title/body rewrite, docker/README.md line 91, the stale Dockerfile comment) are still yours, but the pyproject rationale comment is stale.

[stellarfeline]

@lstein I am working on what requested in your comment. I'll take care of the conflicts.

[stellarfeline]

Thanks for the review @lstein — and for pushing the merge with main, that saved me the conflict resolution. All requested items are now addressed:

1. PR title & description — retitled and rewritten to describe the marker-based approach actually in the diff.
2. Docs — docker/README.md no longer claims x86_64-only support; it now has an "ARM64 / aarch64 (CPU-only)" section with the --build-arg GPU_DRIVER=cpu one-liner and the docker buildx --platform linux/arm64 cross-compile command. The stale Dockerfile comment is refreshed too.
3. Comment accuracy — the pyproject.toml comments now state the real blocker (no aarch64 torchvision wheels on the WHL indexes; torch itself has them).
4. CUDA fallback — the silent CPU fallback for --extra cuda on aarch64 is documented in pyproject.toml, the Dockerfile comment, and the README, with the Grace/GH200 breadcrumb noting that torch+cu128 aarch64 wheels exist and a torchvision-specific workaround would enable real CUDA-on-arm later.

I also re-verified your merge commit on real hardware (Raspberry Pi 5, Ubuntu 24.04, aarch64): uv lock --locked passes, uv sync --frozen --extra cpu installs the PyPI torch 2.7.1 / torchvision 0.22.1 wheels, and both import cleanly. One noteworthy detail: your explicit aarch64 fallback pins in the extras are load-bearing beyond the conflict-universe issue — with the marker-only formulation from my earlier revision, re-locking with current uv (0.11.x) unifies the base torchvision onto the +cpu WHL-index version in the aarch64 fork and breaks the install, while the explicit-pin formulation locks correctly across uv versions. I confirmed the lock in this branch evaluates to PyPI torch/torchvision for all three extras on linux/aarch64.

Happy to do the resolution-only CI guard (uv pip compile --python-platform aarch64-unknown-linux-gnu) as a follow-up PR as suggested.

[stellarfeline]

Follow-up commit a1273f6c02: extended the docs coverage beyond docker, since this change also enables native (non-docker) ARM64 installs — hardware requirements now mention Linux ARM64 CPU-only support, the manual install guide warns not to use --torch-backend=cpu on aarch64 (that index has no aarch64 torchvision wheels; plain PyPI resolution works), the dev environment guide notes --extra cpu for ARM64 uv sync, and the Docker configuration page documents GPU_DRIVER=cpu.

[stellarfeline]

Opened the offered CI-guard follow-up as #9357 (stacked on this PR — its diff reduces to one workflow file once this merges).

[lstein]

All change requests satisfied.
Nice work, thank you for your contribution!