fix: run text encoders on MPS instead of CPU on Apple Silicon

[ChrisLundquist]

Re-lands #12809 (reverted in #13070) with guards that keep fp8/quantized text encoders on the CPU.

Problem

On Apple Silicon vram_state is VRAMState.SHARED, so text_encoder_device() never returns the GPU and text encoders always run on the CPU. For LM-style encoders this dominates generation time: ACE-Step 1.5 on an M5 Max spends 6m19s in LM sampling on CPU vs ~40s on MPS (measured via --gpu-only, which takes the same placement path). Reported for MPS in #12271.

Why #12809 was reverted, and how this addresses it

PyTorch's MPS backend has no float8 support beyond raw storage — any op touching an fp8 MPS tensor, including the .to(fp16) dequantization cast, raises TypeError (pytorch/pytorch#132624, still open, not fixed in nightly). The existing supports_cast() fallback in CLIP.__init__ only inspects declared model dtypes, which never reflect fp8 weights behind comfy_quant metadata (QuantizedTensor reports the compute dtype, not the fp8 storage dtype). So #12809 let quantized text encoders load onto MPS and crash at encode time.

Changes

• add VRAMState.SHARED to text_encoder_device() — SHARED is only ever set on MPS, so other platforms are unaffected
• CLIP.__init__ decides the device before any weights are allocated: it checks supports_cast() on the resolved dtype, and when the load device cannot cast fp8 it scans the incoming state dict for fp8 tensors or comfy_quant markers and keeps those text encoders on the offload device; devices that can cast fp8 (cuda etc.) skip the scan entirely
• the pre-existing "Had to shift TE back" path now updates the current device so the load log stays accurate

Tests

tests-unit/comfy_test/text_encoder_mps_test.py (skipped off-MPS; runs on macOS CI runners): fp16 placement on the GPU, fp8 fallback via declared dtype / secondary dtype (dtype_llama style) / state-dict fp8 weights / comfy_quant markers, and a canary that fails when a torch release adds fp8 casts on MPS so supports_cast() can be relaxed.

Verified with real checkpoints on an M5 Max (torch 2.10.0): the ACE-Step qwen 0.6b+4b text encoders load on mps (CLIP/text encoder model load device: mps); the same LM checkpoint cast to fp8 stays on cpu.

🤖 Generated with Claude Code

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro Plus

Run ID: e22a54ab-9bdc-4af7-b379-62f71030ef52

📥 Commits

Reviewing files that changed from the base of the PR and between c00d1d7 and dc105f8.

📒 Files selected for processing (1)

• tests-unit/comfy_test/text_encoder_mps_test.py

📜 Recent review details

🧰 Additional context used

📓 Path-based instructions (2)

**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

**/*.py: Keep imports at module scope; avoid inline imports unless they are already part of an established optional-backend probe or are needed to avoid an import cycle.
Do not add unnecessary try/except blocks; use them only for optional dependency, platform, or backend capability detection when the program has a useful fallback, and prefer specific exception types when changing new code.
Remove workarounds for PyTorch versions that ComfyUI no longer officially supports; if a workaround does not have a comment naming the exact supported PyTorch version(s), remove it.
Let unsupported model formats, invalid quantization metadata, and bad states fail with clear errors instead of silently producing lower quality output.
Match the existing local style in the file you edit; this codebase tolerates long lines, simple helper functions, module-level state, and direct tensor operations when they make the code easier to follow.
Keep comments sparse and useful; strip useless comments that restate the code or describe obvious behavior, while allowing short TODOs that name the concrete missing follow-up.
Treat dtype, device placement, VRAM usage, and offloading behavior as core correctness concerns; check CPU, CUDA, ROCm, MPS, DirectML, XPU, NPU, and low VRAM implications when touching shared execution or loading code.
Prefer native ComfyUI formats and existing quantization/offload helpers over adding parallel code paths; use comfy.quant_ops, comfy.model_management, comfy.memory_management, comfy.pinned_memory, comfy_aimdo, and comfy-kitchen helpers where they already solve the problem.
Use optimized comfy-kitchen ops where they improve performance without changing expected dtype, device, memory, or interface behavior.
All models should use the optimized attention function selected by ComfyUI; higher-level code must not inspect function identity, names, modules, or implementation details to decide behavior.
Apply the same opacity rule to similar patterns beyond...

Files:

• tests-unit/comfy_test/text_encoder_mps_test.py

**

⚙️ CodeRabbit configuration file

**: ## Engineering Style

• Keep changes small and direct. Most fixes should touch the narrowest code path
  that explains the bug, performance issue, dtype issue, model-format issue, or
  user-facing behavior.
• Change the least amount of files possible. A change that touches many files is
  more likely to be a bad change than a good one unless the broader scope is
  directly required.
• Prefer practical fixes over broad architecture work. Add abstractions only
  when they remove real repeated logic or match an existing ComfyUI pattern.
• Prefer fewer dependencies. Do not add new dependencies to ComfyUI unless they
  are absolutely necessary.
• Delete obsolete code aggressively when newer infrastructure makes it useless.
  Remove dead fallbacks, migration paths, unused options, debug prints, and
  compatibility branches that are no longer needed. Do not leave dead branches,
  unreachable code, or functions that are never called. If code is not
  necessary for the current behavior, remove it.
• Revert or disable problematic behavior quickly when it breaks users. It is
  better to remove a broken feature path than keep a complicated partial fix.
• Preserve existing APIs, node names, model-loading behavior, file layout, and
  workflow compatibility unless the change is explicitly about replacing them.
• Code must look hand-written for this repository. Changes that read like
  generic AI-generated code will be rejected automatically: unnecessary helper
  layers, vague names, boilerplate comments, defensive branches without a real
  failure mode, broad rewrites, or code that ignores the local style.

Architecture Boundaries

• Keep each layer focused on the concepts it owns. Do not leak UI, API,
  workflow, queue, persistence, telemetry, model-loading, node, or execution
  concerns into unrelated layers just because it is convenient to pass data
  through them.
• Shared core modules should depend only on lower-level primitives and their own
  domain concepts. Highe...

Files:

• tests-unit/comfy_test/text_encoder_mps_test.py

---

⚙️ CodeRabbit configuration file

**: IMPORTANT: Only comment on issues directly introduced by this PR's code changes.
Treat AGENTS.md as mandatory repository policy, not optional style guidance.
Flag PR changes that violate AGENTS.md even when the code is otherwise functional.
In particular, enforce architecture boundaries, dtype/device/memory rules,
interface contracts, import style, no unnecessary try/except blocks, no inline
imports, no outbound internet paths in core ComfyUI, and narrow scoped fixes.
Prefer direct findings over suggestions when a rule is violated. Only ignore
AGENTS.md when it clearly conflicts with a newer explicit maintainer instruction
in the PR.
Do NOT flag pre-existing issues in code that was merely moved, re-indented,
de-indented, or reformatted without logic changes. If code appears in the diff
only due to whitespace or structural reformatting (e.g., removing a with: block),
treat it as unchanged. Contributors should not feel obligated to address
pre-existing issues outside the scope of their contribution.

Files:

• tests-unit/comfy_test/text_encoder_mps_test.py

🧠 Learnings (1)

📚 Learning: 2026-02-21T14:01:41.482Z

Learnt from: pythongosssss
Repo: Comfy-Org/ComfyUI PR: 12555
File: comfy_extras/nodes_glsl.py:719-724
Timestamp: 2026-02-21T14:01:41.482Z
Learning: In PyOpenGL, bare Python scalars can be accepted for 1-element array parameters by NumberHandler. This means you can pass an int/float directly to OpenGL texture deletion (e.g., glDeleteTextures(tex)) without wrapping in a list. Verify function-specific expectations and ensure types match what the OpenGL call expects; use explicit lists only when the API requires an array.

Applied to files:

• tests-unit/comfy_test/text_encoder_mps_test.py

🔇 Additional comments (2)

tests-unit/comfy_test/text_encoder_mps_test.py (2)

68-141: LGTM!

---

143-153: 🎯 Functional Correctness

Check the MPS skip scope for these two tests. If the file-wide MPS gate covers them, move test_fp8_capable_devices_skip_the_quant_fallback and test_low_vram_states_keep_text_encoders_on_cpu out or mark them separately so the CUDA/low-VRAM paths run on standard CI.

---

📝 Walkthrough

Walkthrough

Changes

This pull request updates text encoder device selection to include the SHARED VRAM state for GPU placement. It adds fallback-to-CPU handling in CLIP.__init__ when the requested dtype cannot be cast on the chosen load device, and when fp8 or quantized tensors are detected in the state dict without fp8 cast support. When that fallback happens during later dtype checks, the constructor parameters now keep params['device'] aligned with the offload device. A new MPS-focused test module covers device placement and fallback behavior across fp16, fp8, quantized, mixed-dtype, and VRAM-state cases.

Sequence Diagram(s)

Not applicable; the change is mainly conditional device-selection logic plus unit tests.

Related Issues: Not specified.

Related PRs: Not specified.

Suggested labels: bug, tests, mps

Suggested reviewers: Not specified.

Poem
A rabbit hops on silicon shores,
Where fp8 bits knock at MPS doors.
SHARED VRAM joins the GPU crew,
And CPU fallback keeps things true.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: moving text encoders onto MPS on Apple Silicon.
Description check	✅ Passed	The description matches the changeset and explains the MPS placement fix plus fp8/quantized safeguards.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}