Fix partial-load device recovery across CLIP, T5, and Qwen text encoders

invokeai/app/invocations/compel.py

@@ -19,6 +19,7 @@
 from invokeai.app.invocations.primitives import ConditioningOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.app.util.ti_utils import generate_ti_list
+from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device
 from invokeai.backend.model_patcher import ModelPatcher
 from invokeai.backend.patches.layer_patcher import LayerPatcher
 from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
@@ -103,7 +104,7 @@ def _lora_loader() -> Iterator[Tuple[ModelPatchRaw, float]]:
                 textual_inversion_manager=ti_manager,
                 dtype_for_device_getter=TorchDevice.choose_torch_dtype,
                 truncate_long_prompts=False,
-                device=text_encoder.device,  # Use the device the model is actually on
+                device=get_effective_device(text_encoder),
                 split_long_text_mode=SplitLongTextMode.SENTENCES,
             )
 
@@ -212,7 +213,7 @@ def _lora_loader() -> Iterator[Tuple[ModelPatchRaw, float]]:
                 truncate_long_prompts=False,  # TODO:
                 returned_embeddings_type=ReturnedEmbeddingsType.PENULTIMATE_HIDDEN_STATES_NON_NORMALIZED,  # TODO: clip skip
                 requires_pooled=get_pooled,
-                device=text_encoder.device,  # Use the device the model is actually on
+                device=get_effective_device(text_encoder),
                 split_long_text_mode=SplitLongTextMode.SENTENCES,
             )
 

invokeai/app/invocations/sd3_text_encoder.py

@@ -16,6 +16,7 @@
 from invokeai.app.invocations.model import CLIPField, T5EncoderField
 from invokeai.app.invocations.primitives import SD3ConditioningOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
+from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device
 from invokeai.backend.model_manager.taxonomy import ModelFormat
 from invokeai.backend.patches.layer_patcher import LayerPatcher
 from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_CLIP_PREFIX
@@ -103,6 +104,7 @@ def _t5_encode(self, context: InvocationContext, max_seq_len: int) -> torch.Tens
             context.util.signal_progress("Running T5 encoder")
             assert isinstance(t5_text_encoder, T5EncoderModel)
             assert isinstance(t5_tokenizer, (T5Tokenizer, T5TokenizerFast))
+            t5_device = get_effective_device(t5_text_encoder)
 
             text_inputs = t5_tokenizer(
                 prompt,
@@ -125,7 +127,7 @@ def _t5_encode(self, context: InvocationContext, max_seq_len: int) -> torch.Tens
                     f" {max_seq_len} tokens: {removed_text}"
                 )
 
-            prompt_embeds = t5_text_encoder(text_input_ids.to(t5_text_encoder.device))[0]
+            prompt_embeds = t5_text_encoder(text_input_ids.to(t5_device))[0]
 
         assert isinstance(prompt_embeds, torch.Tensor)
         return prompt_embeds
@@ -144,6 +146,7 @@ def _clip_encode(
             context.util.signal_progress("Running CLIP encoder")
             assert isinstance(clip_text_encoder, (CLIPTextModel, CLIPTextModelWithProjection))
             assert isinstance(clip_tokenizer, CLIPTokenizer)
+            clip_device = get_effective_device(clip_text_encoder)
 
             clip_text_encoder_config = clip_text_encoder_info.config
             assert clip_text_encoder_config is not None
@@ -187,9 +190,7 @@ def _clip_encode(
                     "The following part of your input was truncated because CLIP can only handle sequences up to"
                     f" {tokenizer_max_length} tokens: {removed_text}"
                 )
-            prompt_embeds = clip_text_encoder(
-                input_ids=text_input_ids.to(clip_text_encoder.device), output_hidden_states=True
-            )
+            prompt_embeds = clip_text_encoder(input_ids=text_input_ids.to(clip_device), output_hidden_states=True)
             pooled_prompt_embeds = prompt_embeds[0]
             prompt_embeds = prompt_embeds.hidden_states[-2]
 

invokeai/backend/flux/modules/conditioner.py

@@ -3,6 +3,8 @@
 from torch import Tensor, nn
 from transformers import PreTrainedModel, PreTrainedTokenizer, PreTrainedTokenizerFast
 
+from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device
+
 
 class HFEncoder(nn.Module):
     def __init__(
@@ -32,7 +34,7 @@ def forward(self, text: list[str]) -> Tensor:
         )
 
         # Move inputs to the same device as the model to support cpu_only models
-        model_device = next(self.hf_module.parameters()).device
+        model_device = get_effective_device(self.hf_module)
 
         outputs = self.hf_module(
             input_ids=batch_encoding["input_ids"].to(model_device),

invokeai/backend/model_manager/load/load_base.py

@@ -58,7 +58,12 @@ def __init__(self, cache_record: CacheRecord, cache: ModelCache):
 
     def __enter__(self) -> AnyModel:
         self._cache.lock(self._cache_record, None)
-        return self.model
+        try:
+            self.repair_required_tensors_on_device()
+            return self.model
+        except Exception:
+            self._cache.unlock(self._cache_record)
+            raise
 
     def __exit__(self, *args: Any, **kwargs: Any) -> None:
         self._cache.unlock(self._cache_record)
@@ -74,6 +79,7 @@ def model_on_device(
         """
         self._cache.lock(self._cache_record, working_mem_bytes)
         try:
+            self.repair_required_tensors_on_device()
             yield (self._cache_record.cached_model.get_cpu_state_dict(), self._cache_record.cached_model.model)
         finally:
             self._cache.unlock(self._cache_record)

tests/app/invocations/test_compel.py

@@ -0,0 +1,139 @@
+from contextlib import contextmanager, nullcontext
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import torch
+
+from invokeai.app.invocations.compel import SDXLPromptInvocationBase
+
+
+class FakeClipTextEncoder(torch.nn.Module):
+    def __init__(self, effective_device: torch.device):
+        super().__init__()
+        self.register_parameter("cpu_param", torch.nn.Parameter(torch.ones(1)))
+        self.register_buffer("active_buffer", torch.ones(1, device=effective_device))
+        self.dtype = torch.float32
+
+    @property
+    def device(self) -> torch.device:
+        return torch.device("cpu")
+
+
+class FakeTokenizer:
+    pass
+
+
+class FakeLoadedModel:
+    def __init__(self, model, config=None):
+        self._model = model
+        self.config = config
+
+    @contextmanager
+    def model_on_device(self):
+        yield (None, self._model)
+
+    def __enter__(self):
+        return self._model
+
+    def __exit__(self, exc_type, exc, tb):
+        return False
+
+
+class FakeCompel:
+    last_init_device: torch.device | None = None
+
+    def __init__(self, *args, device: torch.device, **kwargs):
+        del args, kwargs
+        FakeCompel.last_init_device = device
+        self.conditioning_provider = SimpleNamespace(
+            get_pooled_embeddings=lambda prompts: torch.ones((len(prompts), 4), dtype=torch.float32)
+        )
+
+    @staticmethod
+    def parse_prompt_string(prompt: str) -> str:
+        return prompt
+
+    def build_conditioning_tensor_for_conjunction(self, conjunction: str):
+        del conjunction
+        return torch.ones((1, 4, 4), dtype=torch.float32), {}
+
+
+@contextmanager
+def fake_apply_ti(tokenizer, text_encoder, ti_list):
+    del text_encoder, ti_list
+    yield tokenizer, object()
+
+
+def test_sdxl_run_clip_compel_uses_effective_device_for_partially_loaded_model(monkeypatch):
+    module_path = "invokeai.app.invocations.compel"
+    effective_device = torch.device("meta")
+    text_encoder = FakeClipTextEncoder(effective_device=effective_device)
+    tokenizer = FakeTokenizer()
+    text_encoder_info = FakeLoadedModel(text_encoder, config=SimpleNamespace(base="sdxl"))
+    tokenizer_info = FakeLoadedModel(tokenizer)
+
+    mock_context = MagicMock()
+    mock_context.models.load.side_effect = [text_encoder_info, tokenizer_info]
+    mock_context.config.get.return_value.log_tokenization = False
+    mock_context.util.signal_progress = MagicMock()
+
+    monkeypatch.setattr(f"{module_path}.CLIPTextModel", FakeClipTextEncoder)
+    monkeypatch.setattr(f"{module_path}.CLIPTextModelWithProjection", FakeClipTextEncoder)
+    monkeypatch.setattr(f"{module_path}.CLIPTokenizer", FakeTokenizer)
+    monkeypatch.setattr(f"{module_path}.Compel", FakeCompel)
+    monkeypatch.setattr(f"{module_path}.generate_ti_list", lambda prompt, base, context: [])
+    monkeypatch.setattr(f"{module_path}.LayerPatcher.apply_smart_model_patches", lambda **kwargs: nullcontext())
+    monkeypatch.setattr(f"{module_path}.ModelPatcher.apply_clip_skip", lambda *args, **kwargs: nullcontext())
+    monkeypatch.setattr(f"{module_path}.ModelPatcher.apply_ti", fake_apply_ti)
+
+    base = SDXLPromptInvocationBase()
+    cond, pooled = base.run_clip_compel(
+        context=mock_context,
+        clip_field=SimpleNamespace(
+            text_encoder=SimpleNamespace(), tokenizer=SimpleNamespace(), loras=[], skipped_layers=0
+        ),
+        prompt="test prompt",
+        get_pooled=False,
+        lora_prefix="lora_te1_",
+        zero_on_empty=False,
+    )
+
+    assert FakeCompel.last_init_device == effective_device
+    assert cond.shape == (1, 4, 4)
+    assert pooled is None
+
+
+def test_sdxl_run_clip_compel_uses_cpu_for_fully_cpu_model(monkeypatch):
+    module_path = "invokeai.app.invocations.compel"
+    text_encoder = FakeClipTextEncoder(effective_device=torch.device("cpu"))
+    tokenizer = FakeTokenizer()
+    text_encoder_info = FakeLoadedModel(text_encoder, config=SimpleNamespace(base="sdxl"))
+    tokenizer_info = FakeLoadedModel(tokenizer)
+
+    mock_context = MagicMock()
+    mock_context.models.load.side_effect = [text_encoder_info, tokenizer_info]
+    mock_context.config.get.return_value.log_tokenization = False
+    mock_context.util.signal_progress = MagicMock()
+
+    monkeypatch.setattr(f"{module_path}.CLIPTextModel", FakeClipTextEncoder)
+    monkeypatch.setattr(f"{module_path}.CLIPTextModelWithProjection", FakeClipTextEncoder)
+    monkeypatch.setattr(f"{module_path}.CLIPTokenizer", FakeTokenizer)
+    monkeypatch.setattr(f"{module_path}.Compel", FakeCompel)
+    monkeypatch.setattr(f"{module_path}.generate_ti_list", lambda prompt, base, context: [])
+    monkeypatch.setattr(f"{module_path}.LayerPatcher.apply_smart_model_patches", lambda **kwargs: nullcontext())
+    monkeypatch.setattr(f"{module_path}.ModelPatcher.apply_clip_skip", lambda *args, **kwargs: nullcontext())
+    monkeypatch.setattr(f"{module_path}.ModelPatcher.apply_ti", fake_apply_ti)
+
+    base = SDXLPromptInvocationBase()
+    base.run_clip_compel(
+        context=mock_context,
+        clip_field=SimpleNamespace(
+            text_encoder=SimpleNamespace(), tokenizer=SimpleNamespace(), loras=[], skipped_layers=0
+        ),
+        prompt="test prompt",
+        get_pooled=False,
+        lora_prefix="lora_te1_",
+        zero_on_empty=False,
+    )
+
+    assert FakeCompel.last_init_device == torch.device("cpu")