diff --git a/invokeai/app/invocations/cogview4_text_encoder.py b/invokeai/app/invocations/cogview4_text_encoder.py index c6ef1663cf8..3b5b1dc73f2 100644 --- a/invokeai/app/invocations/cogview4_text_encoder.py +++ b/invokeai/app/invocations/cogview4_text_encoder.py @@ -6,6 +6,7 @@ from invokeai.app.invocations.model import GlmEncoderField from invokeai.app.invocations.primitives import CogView4ConditioningOutput 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.stable_diffusion.diffusion.conditioning_data import ( CogView4ConditioningInfo, ConditioningFieldData, @@ -46,10 +47,18 @@ def _glm_encode(self, context: InvocationContext, max_seq_len: int) -> torch.Ten prompt = [self.prompt] # TODO(ryand): Add model inputs to the invocation rather than hard-coding. + glm_text_encoder_info = context.models.load(self.glm_encoder.text_encoder) with ( - context.models.load(self.glm_encoder.text_encoder).model_on_device() as (_, glm_text_encoder), + glm_text_encoder_info.model_on_device() as (_, glm_text_encoder), context.models.load(self.glm_encoder.tokenizer).model_on_device() as (_, glm_tokenizer), ): + repaired_tensors = glm_text_encoder_info.repair_required_tensors_on_device() + device = get_effective_device(glm_text_encoder) + if repaired_tensors > 0: + context.logger.warning( + f"Recovered {repaired_tensors} required GLM tensor(s) onto {device} after a partial device mismatch." + ) + context.util.signal_progress("Running GLM text encoder") assert isinstance(glm_text_encoder, GlmModel) assert isinstance(glm_tokenizer, PreTrainedTokenizerFast) @@ -85,9 +94,7 @@ def _glm_encode(self, context: InvocationContext, max_seq_len: int) -> torch.Ten device=text_input_ids.device, ) text_input_ids = torch.cat([pad_ids, text_input_ids], dim=1) - prompt_embeds = glm_text_encoder( - text_input_ids.to(glm_text_encoder.device), output_hidden_states=True - ).hidden_states[-2] + prompt_embeds = glm_text_encoder(text_input_ids.to(device), output_hidden_states=True).hidden_states[-2] assert isinstance(prompt_embeds, torch.Tensor) return prompt_embeds diff --git a/invokeai/app/invocations/flux2_klein_text_encoder.py b/invokeai/app/invocations/flux2_klein_text_encoder.py index 6ca307ebf02..b44e782c8ad 100644 --- a/invokeai/app/invocations/flux2_klein_text_encoder.py +++ b/invokeai/app/invocations/flux2_klein_text_encoder.py @@ -25,6 +25,7 @@ from invokeai.app.invocations.model import Qwen3EncoderField from invokeai.app.invocations.primitives import FluxConditioningOutput 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.patches.layer_patcher import LayerPatcher from invokeai.backend.patches.lora_conversions.flux_lora_constants import FLUX_LORA_T5_PREFIX from invokeai.backend.patches.model_patch_raw import ModelPatchRaw @@ -100,7 +101,12 @@ def _encode_prompt(self, context: InvocationContext, exit_stack: ExitStack) -> T tokenizer_info = context.models.load(self.qwen3_encoder.tokenizer) (_, tokenizer) = exit_stack.enter_context(tokenizer_info.model_on_device()) - device = text_encoder.device + repaired_tensors = text_encoder_info.repair_required_tensors_on_device() + device = get_effective_device(text_encoder) + if repaired_tensors > 0: + context.logger.warning( + f"Recovered {repaired_tensors} required Qwen3 tensor(s) onto {device} after a partial device mismatch." + ) # Apply LoRA models lora_dtype = TorchDevice.choose_bfloat16_safe_dtype(device) diff --git a/invokeai/app/invocations/z_image_text_encoder.py b/invokeai/app/invocations/z_image_text_encoder.py index 06718c4897f..c3405d6dc8b 100644 --- a/invokeai/app/invocations/z_image_text_encoder.py +++ b/invokeai/app/invocations/z_image_text_encoder.py @@ -16,6 +16,7 @@ from invokeai.app.invocations.model import Qwen3EncoderField from invokeai.app.invocations.primitives import ZImageConditioningOutput 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.patches.layer_patcher import LayerPatcher from invokeai.backend.patches.lora_conversions.z_image_lora_constants import Z_IMAGE_LORA_QWEN3_PREFIX from invokeai.backend.patches.model_patch_raw import ModelPatchRaw @@ -76,11 +77,17 @@ def _encode_prompt(self, context: InvocationContext, max_seq_len: int) -> torch. tokenizer_info = context.models.load(self.qwen3_encoder.tokenizer) with ExitStack() as exit_stack: - (_, text_encoder) = exit_stack.enter_context(text_encoder_info.model_on_device()) + (cached_weights, text_encoder) = exit_stack.enter_context(text_encoder_info.model_on_device()) (_, tokenizer) = exit_stack.enter_context(tokenizer_info.model_on_device()) - # Use the device that the text_encoder is actually on - device = text_encoder.device + # Use the device that the text encoder is effectively executing on, and repair any required tensors left on + # the CPU by a previous interrupted run. + repaired_tensors = text_encoder_info.repair_required_tensors_on_device() + device = get_effective_device(text_encoder) + if repaired_tensors > 0: + context.logger.warning( + f"Recovered {repaired_tensors} required Qwen3 tensor(s) onto {device} after a partial device mismatch." + ) # Apply LoRA models to the text encoder lora_dtype = TorchDevice.choose_bfloat16_safe_dtype(device) @@ -90,6 +97,7 @@ def _encode_prompt(self, context: InvocationContext, max_seq_len: int) -> torch. patches=self._lora_iterator(context), prefix=Z_IMAGE_LORA_QWEN3_PREFIX, dtype=lora_dtype, + cached_weights=cached_weights, ) ) diff --git a/invokeai/backend/model_manager/load/load_base.py b/invokeai/backend/model_manager/load/load_base.py index a4004afba75..b972969a687 100644 --- a/invokeai/backend/model_manager/load/load_base.py +++ b/invokeai/backend/model_manager/load/load_base.py @@ -14,6 +14,9 @@ from invokeai.app.services.config import InvokeAIAppConfig from invokeai.backend.model_manager.configs.factory import AnyModelConfig from invokeai.backend.model_manager.load.model_cache.cache_record import CacheRecord +from invokeai.backend.model_manager.load.model_cache.cached_model.cached_model_with_partial_load import ( + CachedModelWithPartialLoad, +) from invokeai.backend.model_manager.load.model_cache.model_cache import ModelCache from invokeai.backend.model_manager.taxonomy import AnyModel, SubModelType @@ -80,6 +83,13 @@ def model(self) -> AnyModel: """Return the model without locking it.""" return self._cache_record.cached_model.model + def repair_required_tensors_on_device(self) -> int: + """Repair required tensors that should be resident on the cached model's execution device.""" + cached_model = self._cache_record.cached_model + if not isinstance(cached_model, CachedModelWithPartialLoad): + return 0 + return cached_model.repair_required_tensors_on_compute_device() + class LoadedModel(LoadedModelWithoutConfig): """Context manager object that mediates transfer from RAM<->VRAM.""" diff --git a/invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_with_partial_load.py b/invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_with_partial_load.py index f80b017ba7b..328978b45b1 100644 --- a/invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_with_partial_load.py +++ b/invokeai/backend/model_manager/load/model_cache/cached_model/cached_model_with_partial_load.py @@ -149,6 +149,27 @@ def full_unload_from_vram(self) -> int: """Unload all weights from VRAM.""" return self.partial_unload_from_vram(self.total_bytes()) + @torch.no_grad() + def repair_required_tensors_on_compute_device(self) -> int: + """Repair required non-autocast tensors that were left off the compute device. + + This can happen if an interrupted run leaves the model in a partially inconsistent state. Any repaired device + movement invalidates the cached VRAM accounting. + """ + cur_state_dict = self._model.state_dict() + keys_to_repair = { + key + for key in self._keys_in_modules_that_do_not_support_autocast + if cur_state_dict[key].device.type != self._compute_device.type + } + if len(keys_to_repair) == 0: + return 0 + + self._load_state_dict_with_device_conversion(cur_state_dict, keys_to_repair, self._compute_device) + self._move_non_persistent_buffers_to_device(self._compute_device) + self._cur_vram_bytes = None + return len(keys_to_repair) + def _load_state_dict_with_device_conversion( self, state_dict: dict[str, torch.Tensor], keys_to_convert: set[str], target_device: torch.device ): diff --git a/tests/app/invocations/test_cogview4_text_encoder.py b/tests/app/invocations/test_cogview4_text_encoder.py new file mode 100644 index 00000000000..81741d41384 --- /dev/null +++ b/tests/app/invocations/test_cogview4_text_encoder.py @@ -0,0 +1,80 @@ +from contextlib import contextmanager +from types import SimpleNamespace +from unittest.mock import MagicMock + +import torch + +from invokeai.app.invocations.cogview4_text_encoder import CogView4TextEncoderInvocation + + +class FakeGlmModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_parameter("weight", torch.nn.Parameter(torch.ones(1))) + self.repaired = False + self.forward_input_device: torch.device | None = None + + def forward(self, input_ids: torch.Tensor, output_hidden_states: bool = False): + assert output_hidden_states + if not self.repaired: + raise RuntimeError("model must be repaired before forward") + + self.forward_input_device = input_ids.device + hidden = input_ids.unsqueeze(-1).float() + return SimpleNamespace(hidden_states=[hidden, hidden + 1]) + + +class FakeTokenizer: + pad_token_id = 0 + + def __call__(self, prompt, padding, max_length=None, truncation=None, add_special_tokens=None, return_tensors=None): + del prompt, padding, max_length, truncation, add_special_tokens, return_tensors + return SimpleNamespace(input_ids=torch.tensor([[1, 2, 3]], dtype=torch.long)) + + def batch_decode(self, input_ids): + del input_ids + return ["decoded"] + + +class FakeLoadedModel: + def __init__(self, model): + self._model = model + self.repair_calls = 0 + + @contextmanager + def model_on_device(self): + yield (None, self._model) + + def repair_required_tensors_on_device(self) -> int: + self.repair_calls += 1 + self._model.repaired = True + return 1 + + +def test_cogview4_text_encoder_repairs_model_before_forward(monkeypatch): + fake_model = FakeGlmModel() + fake_tokenizer = FakeTokenizer() + fake_model_info = FakeLoadedModel(fake_model) + fake_tokenizer_info = FakeLoadedModel(fake_tokenizer) + + mock_context = MagicMock() + mock_context.models.load.side_effect = [fake_model_info, fake_tokenizer_info] + mock_context.util.signal_progress = MagicMock() + mock_context.logger.warning = MagicMock() + + invocation = CogView4TextEncoderInvocation.model_construct( + prompt="test prompt", + glm_encoder=SimpleNamespace(text_encoder=SimpleNamespace(), tokenizer=SimpleNamespace()), + ) + + module_path = "invokeai.app.invocations.cogview4_text_encoder" + monkeypatch.setattr(f"{module_path}.GlmModel", FakeGlmModel) + monkeypatch.setattr(f"{module_path}.PreTrainedTokenizerFast", FakeTokenizer) + + embeds = invocation._glm_encode(mock_context, max_seq_len=16) + + assert fake_model_info.repair_calls == 1 + mock_context.logger.warning.assert_called_once() + mock_context.util.signal_progress.assert_called_once_with("Running GLM text encoder") + assert fake_model.forward_input_device == torch.device("cpu") + assert embeds.shape == (1, 16, 1) diff --git a/tests/backend/model_manager/load/model_cache/cached_model/test_repair_required_tensors.py b/tests/backend/model_manager/load/model_cache/cached_model/test_repair_required_tensors.py new file mode 100644 index 00000000000..306e655a18c --- /dev/null +++ b/tests/backend/model_manager/load/model_cache/cached_model/test_repair_required_tensors.py @@ -0,0 +1,47 @@ +import pytest +import torch + +from invokeai.backend.model_manager.load.model_cache.cached_model.cached_model_with_partial_load import ( + CachedModelWithPartialLoad, +) +from invokeai.backend.model_manager.load.model_cache.torch_module_autocast.torch_module_autocast import ( + apply_custom_layers_to_model, +) + + +class ModelWithRequiredScale(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(4, 4) + self.scale = torch.nn.Parameter(torch.ones(4)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.linear(x) * self.scale + + +@pytest.mark.parametrize( + "device", + [ + pytest.param( + torch.device("cuda"), marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA device") + ), + pytest.param( + torch.device("mps"), + marks=pytest.mark.skipif(not torch.backends.mps.is_available(), reason="requires MPS device"), + ), + ], +) +@pytest.mark.parametrize("keep_ram_copy", [True, False]) +@torch.no_grad() +def test_repair_required_tensors_on_compute_device(device: torch.device, keep_ram_copy: bool): + model = ModelWithRequiredScale() + apply_custom_layers_to_model(model, device_autocasting_enabled=True) + cached_model = CachedModelWithPartialLoad(model=model, compute_device=device, keep_ram_copy=keep_ram_copy) + + cached_model._cur_vram_bytes = 0 + repaired_tensors = cached_model.repair_required_tensors_on_compute_device() + + assert repaired_tensors == 1 + assert cached_model._cur_vram_bytes is None + assert model.scale.device.type == device.type + assert all(param.device.type == "cpu" for param in model.linear.parameters())