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

comfy/model_management.py

@@ -1126,7 +1126,7 @@ def text_encoder_offload_device():
 def text_encoder_device():
     if args.gpu_only:
         return get_torch_device()
-    elif vram_state in (VRAMState.HIGH_VRAM, VRAMState.NORMAL_VRAM) or comfy.memory_management.aimdo_enabled:
+    elif vram_state in (VRAMState.HIGH_VRAM, VRAMState.NORMAL_VRAM, VRAMState.SHARED) or comfy.memory_management.aimdo_enabled:
         if should_use_fp16(prioritize_performance=False):
             return get_torch_device()
         else:

comfy/sd.py

@@ -235,6 +235,18 @@ def __init__(self, target=None, embedding_directory=None, no_init=False, tokeniz
         if dtype is None:
             dtype = model_management.text_encoder_dtype(load_device)
 
+        if not model_management.supports_cast(load_device, dtype):
+            load_device = offload_device
+
+        if load_device != offload_device and not model_management.supports_cast(load_device, torch.float8_e4m3fn):
+            # Quantized state dicts can contain weights in dtypes the declared
+            # model dtypes never reflect (e.g. fp8 layers behind comfy_quant
+            # metadata), which devices like mps cannot cast.
+            for c in (state_dict if isinstance(state_dict, list) else [state_dict]):
+                if any(k.endswith("comfy_quant") or (isinstance(w, torch.Tensor) and w.dtype in (torch.float8_e4m3fn, torch.float8_e5m2)) for k, w in c.items()):
+                    load_device = offload_device
+                    break
+
         params['dtype'] = dtype
         params['device'] = model_options.get("initial_device", model_management.text_encoder_initial_device(load_device, offload_device, parameters * model_management.dtype_size(dtype)))
         params['model_options'] = model_options
@@ -246,6 +258,7 @@ def __init__(self, target=None, embedding_directory=None, no_init=False, tokeniz
                 load_device = offload_device
                 if params['device'] != offload_device:
                     self.cond_stage_model.to(offload_device)
+                    params['device'] = offload_device
                     logging.warning("Had to shift TE back.")
 
         model_management.archive_model_dtypes(self.cond_stage_model)

tests-unit/comfy_test/text_encoder_mps_test.py

@@ -0,0 +1,153 @@
+"""Text encoder device placement on Apple Silicon (MPS).
+
+MPS machines run in VRAMState.SHARED. Text encoders the device supports
+(fp16/bf16/fp32) should load on the GPU, while fp8/quantized ones must
+stay on the CPU because MPS cannot cast float8 dtypes.
+"""
+
+import pytest
+import torch
+
+import comfy.model_management as mm
+import comfy.ops
+import comfy.sd
+
+mps_only = pytest.mark.skipif(
+    not (torch.backends.mps.is_available() and mm.is_device_mps(mm.get_torch_device())),
+    reason="requires an Apple Silicon MPS device",
+)
+
+FP8_DTYPES = [torch.float8_e4m3fn, torch.float8_e5m2]
+
+# Big enough that text_encoder_initial_device() picks the load device.
+LARGE_PARAM_COUNT = 2 * 1024 * 1024 * 1024
+
+
+class DummyTokenizer:
+    def __init__(self, embedding_directory=None, tokenizer_data={}):
+        pass
+
+
+class DummyTEModel(torch.nn.Module):
+    def __init__(self, device=None, dtype=None, model_options={}):
+        super().__init__()
+        operations = model_options.get("custom_operations", comfy.ops.manual_cast)
+        self.linear = operations.Linear(8, 8, dtype=dtype, device=device)
+        self.dtypes = set([dtype])
+        self.construct_device = torch.device(device) if device is not None else None
+
+    def load_sd(self, sd):
+        return ([], [])
+
+    def forward(self, x):
+        return self.linear(x)
+
+
+def make_clip(dtype, state_dict=[], clip_class=DummyTEModel):
+    class Target:
+        params = {}
+        clip = clip_class
+        tokenizer = DummyTokenizer
+
+    return comfy.sd.CLIP(
+        target=Target(),
+        parameters=LARGE_PARAM_COUNT,
+        model_options={"dtype": dtype},
+        state_dict=state_dict,
+        disable_dynamic=True,
+    )
+
+
+@pytest.fixture(autouse=True)
+def unload_models():
+    yield
+    mm.unload_all_models()
+
+
+@mps_only
+class TestTextEncoderDeviceMPS:
+    def test_vram_state_is_shared(self):
+        assert mm.vram_state == mm.VRAMState.SHARED
+
+    def test_text_encoder_device_is_mps(self):
+        assert mm.is_device_mps(mm.text_encoder_device())
+
+    @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+    def test_castable_dtype_clip_loads_on_mps(self, dtype):
+        clip = make_clip(dtype)
+        assert mm.is_device_mps(clip.patcher.load_device)
+        weight = clip.cond_stage_model.linear.weight
+        assert weight.device.type == "mps"
+        x = torch.ones((1, 8), dtype=dtype, device=weight.device)
+        assert clip.cond_stage_model(x).device.type == "mps"
+
+    @pytest.mark.parametrize("fp8_dtype", FP8_DTYPES)
+    def test_fp8_clip_falls_back_to_cpu(self, fp8_dtype):
+        clip = make_clip(fp8_dtype)
+        assert clip.patcher.load_device == clip.patcher.offload_device
+        assert clip.cond_stage_model.construct_device.type == "cpu"
+        weight = clip.cond_stage_model.linear.weight
+        assert weight.device.type == "cpu"
+        x = torch.ones((1, 8), dtype=torch.float16, device=weight.device)
+        assert clip.cond_stage_model(x).device.type == "cpu"
+
+    @pytest.mark.parametrize("fp8_dtype", FP8_DTYPES)
+    def test_fp8_state_dict_falls_back_to_cpu(self, fp8_dtype):
+        # fp8 weights in the state dict aren't reflected in the declared
+        # model dtypes, e.g. quantized checkpoints with fp16 norm layers.
+        sd = {"linear.weight": torch.zeros((8, 8), dtype=fp8_dtype)}
+        clip = make_clip(torch.float16, state_dict=[sd])
+        assert not mm.is_device_mps(clip.patcher.load_device)
+        assert clip.cond_stage_model.construct_device.type == "cpu"
+
+    def test_comfy_quant_state_dict_falls_back_to_cpu(self):
+        sd = {
+            "linear.weight": torch.zeros((8, 8), dtype=torch.uint8),
+            "linear.comfy_quant": torch.zeros(16, dtype=torch.uint8),
+            "spiece_model": b"not a tensor",
+        }
+        clip = make_clip(torch.float16, state_dict=[sd])
+        assert not mm.is_device_mps(clip.patcher.load_device)
+        assert clip.cond_stage_model.construct_device.type == "cpu"
+
+    def test_fp8_full_model_state_dict_falls_back_to_cpu(self):
+        # Full checkpoints pass a single dict instead of a list.
+        sd = {"linear.weight": torch.zeros((8, 8), dtype=torch.float8_e4m3fn)}
+        clip = make_clip(torch.float16, state_dict=sd)
+        assert not mm.is_device_mps(clip.patcher.load_device)
+        assert clip.cond_stage_model.construct_device.type == "cpu"
+
+    @pytest.mark.parametrize("fp8_dtype", FP8_DTYPES)
+    def test_mixed_declared_dtypes_fall_back_to_cpu(self, fp8_dtype):
+        # A secondary declared dtype (e.g. dtype_llama) can be fp8 while the
+        # primary dtype is fp16.
+        class MixedDtypeTEModel(DummyTEModel):
+            def __init__(self, device=None, dtype=None, model_options={}):
+                super().__init__(device=device, dtype=dtype, model_options=model_options)
+                self.dtypes = set([dtype, fp8_dtype])
+
+        clip = make_clip(torch.float16, clip_class=MixedDtypeTEModel)
+        assert not mm.is_device_mps(clip.patcher.load_device)
+        assert clip.cond_stage_model.linear.weight.device.type == "cpu"
+
+    @pytest.mark.parametrize("fp8_dtype", FP8_DTYPES)
+    def test_fp8_cast_still_unsupported_on_mps(self, fp8_dtype):
+        # If a torch release adds fp8 casts on MPS, supports_cast() can be
+        # updated to let fp8 text encoders onto the GPU (pytorch#132624).
+        assert not mm.supports_cast(mm.get_torch_device(), fp8_dtype)
+        t = torch.zeros(4, dtype=fp8_dtype, device="mps")
+        with pytest.raises((RuntimeError, TypeError)):
+            t.to(torch.float16)
+
+
+@pytest.mark.parametrize("fp8_dtype", FP8_DTYPES)
+def test_fp8_capable_devices_skip_the_quant_fallback(fp8_dtype):
+    # The state dict scan in CLIP.__init__ is gated on this, so devices
+    # that can cast fp8 keep loading quantized text encoders on the GPU.
+    assert mm.supports_cast(torch.device("cuda"), fp8_dtype)
+
+
+@pytest.mark.parametrize("state", [mm.VRAMState.LOW_VRAM, mm.VRAMState.NO_VRAM])
+def test_low_vram_states_keep_text_encoders_on_cpu(monkeypatch, state):
+    monkeypatch.setattr(mm, "vram_state", state)
+    assert mm.text_encoder_device() == torch.device("cpu")