Fix: compute device partial load (invoke-ai#9375)

* fix(vae): decode on the model's intended compute device, not current residency (invoke-ai#9373)

VAE decode inferred its device via get_effective_device(vae). Under partial
loading with VRAM pressure, all VAE weights can be temporarily offloaded to
RAM, so that returned CPU — placing the latents on CPU and, because the
autocast layers follow the input device, running the entire decode on CPU.

Expose LoadedModel.compute_device (the model's intended device, stable across
partial-load residency and still honoring cpu_only) and use it in all 8 decode
invocations. Adds a regression test reproducing the offloaded-weights case.

* fix(text-encoder): encode on the intended compute device, not current residency

Same class of bug as invoke-ai#9373, for text encoders: get_effective_device() returns
CPU when partial loading has offloaded all weights, running the whole encode on
CPU. Fully autocast-capable encoders (e.g. CLIP) are affected because
repair_required_tensors_on_device() pins nothing to the compute device.

Use LoadedModel.compute_device in compel, sd3, z_image, flux2_klein, cogview4
and qwen_image encoders, and thread it through HFEncoder for the FLUX encoder.
Updates the affected tests to assert the intended-device behavior.

* fix(tests): set vae_info.compute_device in qwen/z-image decode mocks

Like the anima decode test, the qwen-image and z-image working-memory
tests build vae_info as a bare MagicMock and drive invoke(). Since invoke-ai#9373
places latents on vae_info.compute_device, latents.to(device=...) raised
TypeError — but both tests wrapped invoke() in `except Exception: pass`,
so the failure was silently swallowed and the decode path never actually
ran. Set compute_device to torch.device("cpu") so the tests exercise the
real decode path instead of masking the error.

* fix(text-encoder): anima encodes on the intended compute device, not current residency

The anima text encoder inferred its device via text_encoder.device (HF
PreTrainedModel residency), which returns CPU when partial loading has
temporarily offloaded all Qwen3 weights to RAM — running the whole encode on
the CPU. This is the same class of bug as invoke-ai#9373; the other text encoders and
VAE decodes in this PR were already fixed, but the anima encoder was missed.

Use LoadedModel.compute_device instead. This also corrects the device passed to
TorchDevice.choose_anima_inference_dtype(). Adds a regression test covering both
the offloaded-accelerator case and the cpu_only case.

---------

Co-authored-by: Lincoln Stein <lincoln.stein@gmail.com>

Author: Pfannkuchensack

invokeai/app/invocations/anima_latents_to_image.py

@@ -28,7 +28,6 @@
 from invokeai.app.invocations.primitives import ImageOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.backend.flux.modules.autoencoder import AutoEncoder as FluxAutoEncoder
-from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device
 from invokeai.backend.util.devices import TorchDevice
 from invokeai.backend.util.vae_working_memory import (
     estimate_vae_working_memory_anima,
@@ -127,8 +126,10 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
                 raise TypeError(f"Expected AutoencoderKLWan or FluxAutoEncoder, got {type(vae).__name__}.")
 
             vae_dtype = next(iter(vae.parameters())).dtype
-            # Use the VAE's actual device (may be CPU if the model is configured cpu_only).
-            latents = latents.to(device=get_effective_device(vae), dtype=vae_dtype)
+            # Use the VAE's intended compute device (CUDA/MPS, or CPU if configured cpu_only). Do NOT infer it from
+            # current param residency: partial loading may have temporarily offloaded all weights to RAM, which would
+            # wrongly place the latents (and thus the whole decode) on the CPU (see #9373).
+            latents = latents.to(device=vae_info.compute_device, dtype=vae_dtype)
 
             TorchDevice.empty_cache()
 

invokeai/app/invocations/anima_text_encoder.py

@@ -125,7 +125,10 @@ def _encode_prompt(
             (_, text_encoder) = exit_stack.enter_context(text_encoder_info.model_on_device())
             (_, tokenizer) = exit_stack.enter_context(tokenizer_info.model_on_device())
 
-            device = text_encoder.device
+            # Use the encoder's intended compute device, not its current parameter residency: partial loading may
+            # have temporarily offloaded all weights to RAM, which would wrongly run the whole encode on the CPU (see
+            # #9373). Qwen3 is fully autocast-capable, so nothing pins it to the compute device otherwise.
+            device = text_encoder_info.compute_device
 
             # Apply LoRA models to the text encoder
             lora_dtype = TorchDevice.choose_anima_inference_dtype(device)

invokeai/app/invocations/cogview4_latents_to_image.py

@@ -17,7 +17,6 @@
 from invokeai.app.invocations.model import VAEField
 from invokeai.app.invocations.primitives import ImageOutput
 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.extensions.seamless import SeamlessExt
 from invokeai.backend.util.devices import TorchDevice
 from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_cogview4
@@ -55,8 +54,10 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
         ):
             context.util.signal_progress("Running VAE")
             assert isinstance(vae, (AutoencoderKL))
-            # Use the VAE's actual device (may be CPU if the model is configured cpu_only).
-            latents = latents.to(get_effective_device(vae))
+            # Use the VAE's intended compute device (CUDA/MPS, or CPU if configured cpu_only). Do NOT infer it from
+            # current param residency: partial loading may have temporarily offloaded all weights to RAM, which would
+            # wrongly place the latents (and thus the whole decode) on the CPU (see #9373).
+            latents = latents.to(vae_info.compute_device)
 
             vae.disable_tiling()
 

invokeai/app/invocations/cogview4_text_encoder.py

@@ -6,7 +6,6 @@
 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,
@@ -52,8 +51,11 @@ def _glm_encode(self, context: InvocationContext, max_seq_len: int) -> torch.Ten
             glm_text_encoder_info.model_on_device() as (_, glm_text_encoder),
             context.models.load(self.glm_encoder.tokenizer).model_on_device() as (_, glm_tokenizer),
         ):
+            # Repair any required tensors left on the CPU by a previous interrupted run, then run on the encoder's
+            # intended compute device. Do NOT infer the device from current parameter residency: partial loading may
+            # have temporarily offloaded all weights to RAM, which would wrongly run the whole encode on the CPU.
             repaired_tensors = glm_text_encoder_info.repair_required_tensors_on_device()
-            device = get_effective_device(glm_text_encoder)
+            device = glm_text_encoder_info.compute_device
             if repaired_tensors > 0:
                 context.logger.warning(
                     f"Recovered {repaired_tensors} required GLM tensor(s) onto {device} after a partial device mismatch."

invokeai/app/invocations/compel.py

@@ -19,7 +19,6 @@
 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
@@ -104,7 +103,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=get_effective_device(text_encoder),
+                device=text_encoder_info.compute_device,
                 split_long_text_mode=SplitLongTextMode.SENTENCES,
             )
 
@@ -213,7 +212,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=get_effective_device(text_encoder),
+                device=text_encoder_info.compute_device,
                 split_long_text_mode=SplitLongTextMode.SENTENCES,
             )
 

invokeai/app/invocations/flux2_klein_text_encoder.py

@@ -25,7 +25,6 @@
 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
@@ -101,8 +100,11 @@ 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())
 
+        # Repair any required tensors left on the CPU by a previous interrupted run, then run on the encoder's
+        # intended compute device. Do NOT infer the device from current parameter residency: partial loading may
+        # have temporarily offloaded all weights to RAM, which would wrongly run the whole encode on the CPU.
         repaired_tensors = text_encoder_info.repair_required_tensors_on_device()
-        device = get_effective_device(text_encoder)
+        device = text_encoder_info.compute_device
         if repaired_tensors > 0:
             context.logger.warning(
                 f"Recovered {repaired_tensors} required Qwen3 tensor(s) onto {device} after a partial device mismatch."

invokeai/app/invocations/flux2_vae_decode.py

@@ -20,7 +20,6 @@
 from invokeai.app.invocations.primitives import ImageOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.backend.model_manager.load.load_base import LoadedModel
-from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device
 from invokeai.backend.util.devices import TorchDevice
 
 
@@ -52,8 +51,10 @@ def _vae_decode(self, vae_info: LoadedModel, latents: torch.Tensor) -> Image.Ima
         """
         with vae_info.model_on_device() as (_, vae):
             vae_dtype = next(iter(vae.parameters())).dtype
-            # Use the VAE's actual device (may be CPU if the model is configured cpu_only).
-            device = get_effective_device(vae)
+            # Use the VAE's intended compute device (CUDA/MPS, or CPU if configured cpu_only). Do NOT infer it from
+            # current param residency: partial loading may have temporarily offloaded all weights to RAM, which would
+            # wrongly place the latents (and thus the whole decode) on the CPU (see #9373).
+            device = vae_info.compute_device
             latents = latents.to(device=device, dtype=vae_dtype)
 
             # Decode using diffusers API

invokeai/app/invocations/flux_text_encoder.py

@@ -115,7 +115,9 @@ def _t5_encode(self, context: InvocationContext) -> torch.Tensor:
                 )
             )
 
-            t5_encoder = HFEncoder(t5_text_encoder, t5_tokenizer, False, self.t5_max_seq_len)
+            t5_encoder = HFEncoder(
+                t5_text_encoder, t5_tokenizer, False, self.t5_max_seq_len, device=t5_encoder_info.compute_device
+            )
 
             if context.config.get().log_tokenization:
                 self._log_t5_tokenization(context, t5_tokenizer)
@@ -158,7 +160,9 @@ def _clip_encode(self, context: InvocationContext) -> torch.Tensor:
                 # There are currently no supported CLIP quantized models. Add support here if needed.
                 raise ValueError(f"Unsupported model format: {clip_text_encoder_config.format}")
 
-            clip_encoder = HFEncoder(clip_text_encoder, clip_tokenizer, True, 77)
+            clip_encoder = HFEncoder(
+                clip_text_encoder, clip_tokenizer, True, 77, device=clip_text_encoder_info.compute_device
+            )
 
             if context.config.get().log_tokenization:
                 self._log_clip_tokenization(context, clip_tokenizer)

invokeai/app/invocations/flux_vae_decode.py

@@ -16,7 +16,6 @@
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.backend.flux.modules.autoencoder import AutoEncoder
 from invokeai.backend.model_manager.load.load_base import LoadedModel
-from invokeai.backend.model_manager.load.model_cache.utils import get_effective_device
 from invokeai.backend.util.devices import TorchDevice
 from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_flux
 
@@ -48,8 +47,10 @@ def _vae_decode(self, vae_info: LoadedModel, latents: torch.Tensor) -> Image.Ima
         with vae_info.model_on_device(working_mem_bytes=estimated_working_memory) as (_, vae):
             assert isinstance(vae, AutoEncoder)
             vae_dtype = next(iter(vae.parameters())).dtype
-            # Use the VAE's actual device (may be CPU if the model is configured cpu_only).
-            latents = latents.to(device=get_effective_device(vae), dtype=vae_dtype)
+            # Use the VAE's intended compute device (CUDA/MPS, or CPU if configured cpu_only). Do NOT infer it from
+            # current param residency: partial loading may have temporarily offloaded all weights to RAM, which would
+            # wrongly place the latents (and thus the whole decode) on the CPU (see #9373).
+            latents = latents.to(device=vae_info.compute_device, dtype=vae_dtype)
             img = vae.decode(latents)
 
         img = img.clamp(-1, 1)

invokeai/app/invocations/latents_to_image.py

@@ -18,7 +18,6 @@
 from invokeai.app.invocations.model import VAEField
 from invokeai.app.invocations.primitives import ImageOutput
 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.extensions.seamless import SeamlessExt
 from invokeai.backend.stable_diffusion.vae_tiling import patch_vae_tiling_params
 from invokeai.backend.util.devices import TorchDevice
@@ -70,8 +69,10 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
         ):
             context.util.signal_progress("Running VAE decoder")
             assert isinstance(vae, (AutoencoderKL, AutoencoderTiny))
-            # Use the VAE's actual device (may be CPU if the model is configured cpu_only).
-            device = get_effective_device(vae)
+            # Use the VAE's intended compute device (CUDA/MPS, or CPU if configured cpu_only). Do NOT infer it from
+            # current param residency: partial loading may have temporarily offloaded all weights to RAM, which would
+            # wrongly place the latents (and thus the whole decode) on the CPU (see #9373).
+            device = vae_info.compute_device
             latents = latents.to(device)
             # Force fp32 when running on CPU (e.g. when the VAE is configured cpu_only). fp16 conv does run on CPU with
             # the pinned torch, but it's much slower than fp32 there, and SD/SDXL VAE has known fp16 overflow issues