fix(vae): honor the loaded model compute device

invokeai/app/invocations/flux/flux_denoise.py

@@ -357,6 +357,8 @@ def _run_diffusion(
 
             if self.control_lora is not None:
                 img_cond = self._prep_structural_control_img_cond(context)
+                if img_cond is not None:
+                    img_cond = img_cond.to(device=device, dtype=inference_dtype)
 
         inpaint_mask = self._prep_inpaint_mask(context, x)
 
@@ -891,7 +893,6 @@ def _prep_flux_fill_img_cond(
         cond_img = np.array(cond_img)
         cond_img = torch.from_numpy(cond_img).float() / 127.5 - 1.0
         cond_img = einops.rearrange(cond_img, "h w c -> 1 c h w")
-        cond_img = cond_img.to(device=device, dtype=dtype)
 
         # Load the mask and resize it to the target image size.
         mask = context.tensors.load(self.fill_conditioning.mask.tensor_name)
@@ -900,14 +901,16 @@ def _prep_flux_fill_img_cond(
         assert mask.dim() == 3
         assert mask.shape[0] == 1
         mask = tv_resize(mask, size=[self.height, self.width], interpolation=tv_transforms.InterpolationMode.NEAREST)
-        mask = mask.to(device=device, dtype=dtype)
+        # Mask on the pixels' device before encoding; a cpu_only VAE need not share the session device.
+        mask = mask.to(device=cond_img.device, dtype=cond_img.dtype)
         mask = einops.rearrange(mask, "1 h w -> 1 1 h w")
 
         # Prepare image conditioning.
         cond_img = cond_img * (1 - mask)
         vae_info = context.models.load(self.controlnet_vae.vae)
         cond_img = FluxVaeEncodeInvocation.vae_encode(vae_info=vae_info, image_tensor=cond_img)
         cond_img = pack(cond_img)
+        mask = mask.to(cond_img)
 
         # Prepare mask conditioning.
         mask = mask[:, 0, :, :]
@@ -916,7 +919,7 @@ def _prep_flux_fill_img_cond(
         mask = pack(mask)
 
         # Merge image and mask conditioning.
-        img_cond = torch.cat((cond_img, mask), dim=-1)
+        img_cond = torch.cat((cond_img, mask), dim=-1).to(device=device, dtype=dtype)
         return img_cond
 
     def _normalize_ip_adapter_fields(self) -> list[IPAdapterField]:

invokeai/app/invocations/vae/anima_image_to_latents.py

@@ -27,7 +27,6 @@
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.backend.model_manager.load.load_base import LoadedModel
 from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor
-from invokeai.backend.util.devices import TorchDevice
 from invokeai.backend.util.qwen_image_vae import (
     QWEN_IMAGE_VAE_MIN_TILE_SIZE,
     as_qwen_image_vae,
@@ -92,7 +91,7 @@ def vae_encode(
             vae = as_qwen_image_vae(vae)
 
             vae_dtype = next(iter(vae.parameters())).dtype
-            image_tensor = image_tensor.to(device=TorchDevice.choose_torch_device(), dtype=vae_dtype)
+            image_tensor = image_tensor.to(device=vae_info.compute_device, dtype=vae_dtype)
 
             # The cached VAE instance is shared with the decode invocation, which may have enabled
             # tiling. That used to be handled by switching tiling off here and leaving it off, which

invokeai/app/invocations/vae/cogview4_image_to_latents.py

@@ -16,7 +16,6 @@
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.backend.model_manager.load.load_base import LoadedModel
 from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor
-from invokeai.backend.util.devices import TorchDevice
 from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_cogview4
 
 # TODO(ryand): This is effectively a copy of SD3ImageToLatentsInvocation and a subset of ImageToLatentsInvocation. We
@@ -48,7 +47,7 @@ def vae_encode(vae_info: LoadedModel, image_tensor: torch.Tensor) -> torch.Tenso
 
             vae.disable_tiling()
 
-            image_tensor = image_tensor.to(device=TorchDevice.choose_torch_device(), dtype=vae.dtype)
+            image_tensor = image_tensor.to(device=vae_info.compute_device, dtype=vae.dtype)
             with torch.inference_mode():
                 image_tensor_dist = vae.encode(image_tensor).latent_dist
                 # TODO: Use seed to make sampling reproducible.

invokeai/app/invocations/vae/ernie_image_vae_decode.py

@@ -23,7 +23,6 @@
 from invokeai.app.invocations.primitives import ImageOutput
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.backend.ernie_image.sampling_utils import unpatchify_latents, vae_denormalize
-from invokeai.backend.util.devices import TorchDevice
 
 
 @invocation(
@@ -46,7 +45,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
 
         vae_info = context.models.load(self.vae.vae)
         with vae_info.model_on_device() as (_, vae):
-            device = TorchDevice.choose_torch_device()
+            device = vae_info.compute_device
             vae_dtype = next(iter(vae.parameters())).dtype
             latents = latents.to(device=device, dtype=vae_dtype)
 

invokeai/app/invocations/vae/qwen_image_image_to_latents.py

@@ -20,7 +20,6 @@
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.backend.model_manager.load.load_base import LoadedModel
 from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor
-from invokeai.backend.util.devices import TorchDevice
 from invokeai.backend.util.qwen_image_vae import (
     QWEN_IMAGE_VAE_MIN_TILE_SIZE,
     as_qwen_image_vae,
@@ -129,7 +128,7 @@ def vae_encode(
             # Reinterpret an Anima-classified Wan VAE as AutoencoderKLQwenImage (identical weights).
             vae = as_qwen_image_vae(vae)
 
-            image_tensor = image_tensor.to(device=TorchDevice.choose_torch_device(), dtype=vae.dtype)
+            image_tensor = image_tensor.to(device=vae_info.compute_device, dtype=vae.dtype)
 
             # Tiling bounds the encode's peak memory to a single tile, which is what makes large
             # inputs (e.g. a 2560x1440 upscale round-trip) encodable while a multi-GB transformer

invokeai/app/invocations/vae/sd3_image_to_latents.py

@@ -16,7 +16,6 @@
 from invokeai.app.services.shared.invocation_context import InvocationContext
 from invokeai.backend.model_manager.load.load_base import LoadedModel
 from invokeai.backend.stable_diffusion.diffusers_pipeline import image_resized_to_grid_as_tensor
-from invokeai.backend.util.devices import TorchDevice
 from invokeai.backend.util.vae_working_memory import estimate_vae_working_memory_sd3
 
 
@@ -44,7 +43,7 @@ def vae_encode(vae_info: LoadedModel, image_tensor: torch.Tensor) -> torch.Tenso
 
             vae.disable_tiling()
 
-            image_tensor = image_tensor.to(device=TorchDevice.choose_torch_device(), dtype=vae.dtype)
+            image_tensor = image_tensor.to(device=vae_info.compute_device, dtype=vae.dtype)
             with torch.inference_mode():
                 image_tensor_dist = vae.encode(image_tensor).latent_dist
                 # TODO: Use seed to make sampling reproducible.

invokeai/backend/flux/extensions/kontext_extension.py

@@ -113,6 +113,7 @@ def _prepare_kontext(self) -> tuple[torch.Tensor, torch.Tensor]:
         canvas_w = 0  # Running canvas width
 
         vae_info = self._context.models.load(self._vae_field.vae)
+        vae_device = vae_info.compute_device
 
         for idx, kontext_field in enumerate(self.kontext_conditioning):
             image = self._context.images.get_pil(kontext_field.image.image_name)
@@ -130,18 +131,17 @@ def _prepare_kontext(self) -> tuple[torch.Tensor, torch.Tensor]:
             # Convert from [0, 1] to [-1, 1] range expected by VAE
             image_tensor = image_tensor * 2.0 - 1.0
             image_tensor = image_tensor.unsqueeze(0)  # Add batch dimension
-            image_tensor = image_tensor.to(self._device)
 
             # Continue with VAE encoding
             # Don't sample from the distribution for reference images - use the mean (matching ComfyUI)
             estimated_working_memory = estimate_vae_working_memory_flux(
-                operation="encode", image_tensor=image_tensor, vae=vae_info.model
+                operation="encode", image_tensor=image_tensor, vae=vae_info.model, device=vae_device
             )
 
             with vae_info.model_on_device(working_memory=estimated_working_memory) as (_, vae):
                 assert isinstance(vae, AutoencoderKL)
                 vae_dtype = next(iter(vae.parameters())).dtype
-                image_tensor = image_tensor.to(device=TorchDevice.choose_torch_device(), dtype=vae_dtype)
+                image_tensor = image_tensor.to(device=vae_device, dtype=vae_dtype)
 
                 def encode_reference(x: torch.Tensor = image_tensor) -> torch.Tensor:
                     # `.mode()` is the distribution mean without noise -- `sample=False` on the

invokeai/backend/flux2/ref_image_extension.py

@@ -119,7 +119,7 @@ def encode_reference_image(
     image_tensor = T.ToTensor()(image) * 2.0 - 1.0
     image_tensor = image_tensor.unsqueeze(0)
 
-    device = TorchDevice.choose_torch_device()
+    device = vae_info.compute_device
     tile_size = reference_encode_tile_size(image_tensor, vae_info.model, device, tiling)
 
     # A tiled encode peaks at one tile; tell the cache up front so it frees the room instead of hitting the

invokeai/backend/util/oom.py

@@ -10,7 +10,8 @@ def is_oom_error(e: RuntimeError) -> bool:
     cuDNN/cuBLAS kernel (e.g. workspace allocation in a VAE's convolutions) arrives as a plain
     RuntimeError, which must be matched by message. XPU exhaustion likewise arrives as a plain
     RuntimeError, naming the Level Zero/UR result code (`..._OUT_OF_DEVICE_MEMORY`) rather than the
-    words "out of memory" -- so it needs its own spelling to be matched here.
+    words "out of memory" -- so it needs its own spelling to be matched here. PyTorch CPU
+    allocations also raise plain RuntimeError with a backend-specific DefaultCPUAllocator message.
 
     `out_of_host_memory` is knowingly over-broad: Level Zero returns it for driver-side resource
     failures generally (kernel compilation, handle exhaustion), not only host allocation. Matching
@@ -31,4 +32,6 @@ def is_oom_error(e: RuntimeError) -> bool:
         or "out_of_host_memory" in msg
         or "cudnn_status_alloc_failed" in msg
         or "cublas_status_alloc_failed" in msg
+        or "defaultcpuallocator: not enough memory: you tried to allocate " in msg
+        or "defaultcpuallocator: can't allocate memory: you tried to allocate " in msg
     )

tests/app/invocations/test_anima_vae_tiling.py

@@ -9,6 +9,7 @@
 
 from unittest.mock import MagicMock, patch
 
+import pytest
 import torch
 from diffusers.models.autoencoders import AutoencoderKLWan
 
@@ -18,6 +19,7 @@
     ANIMA_VAE_TILE_STRIDE,
     AnimaLatentsToImageInvocation,
 )
+from invokeai.app.invocations.vae.qwen_image_image_to_latents import QwenImageImageToLatentsInvocation
 from invokeai.backend.util.devices import TorchDevice
 
 
@@ -133,6 +135,7 @@ def spy_decode(*args, **kwargs):
 def _build_encode_context(vae: AutoencoderKLWan):
     vae_info = MagicMock()
     vae_info.model = vae
+    vae_info.compute_device = torch.device("cpu")
     cm = MagicMock()
     cm.__enter__ = MagicMock(return_value=(None, vae))
     cm.__exit__ = MagicMock(return_value=None)
@@ -146,6 +149,33 @@ def _build_encode_context(vae: AutoencoderKLWan):
     return vae_info, context
 
 
+@pytest.mark.parametrize(
+    ("invocation", "shape"),
+    [
+        (AnimaImageToLatentsInvocation, (1, 16, 4, 4)),
+        (QwenImageImageToLatentsInvocation, (1, 16, 1, 4, 4)),
+    ],
+    ids=["anima", "qwen-image"],
+)
+def test_wan_layout_encode_uses_the_cpu_vae_in_an_accelerator_session(monkeypatch, invocation, shape):
+    vae = _build_tiny_vae()
+    vae_info, _ = _build_encode_context(vae)
+    real_encode = vae.encode
+
+    def encode(pixels, **kwargs):
+        assert pixels.device == torch.device("cpu")
+        assert pixels.dtype == vae.dtype
+        return real_encode(pixels, **kwargs)
+
+    monkeypatch.setattr(vae, "encode", encode)
+    monkeypatch.setattr(TorchDevice, "choose_torch_device", lambda: torch.device("meta"))
+
+    latents = invocation.vae_encode(vae_info, torch.zeros(1, 3, 32, 32, dtype=torch.float64))
+
+    assert latents.shape == shape
+    assert latents.device == torch.device("cpu")
+
+
 def _encode(context, vae, image: torch.Tensor, tiled: bool = False, tile_size: int = 0):
     """Run the encode node end to end, recording the tiling state the encode actually saw."""
     seen: list[tuple] = []

tests/app/invocations/test_flux_cpu_vae_conditioning.py

@@ -0,0 +1,138 @@
+"""CPU VAE conditioning reaches the session's real denoise concatenation boundary.
+
+Meta stands in for a different session device, not for GPU execution. The VAE encodes real
+pixels in fp32 on CPU; only the expensive transformer's forward is stopped at its input.
+"""
+
+from contextlib import nullcontext
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+from diffusers import FlowMatchEulerDiscreteScheduler
+from PIL import Image
+
+from invokeai.app.invocations.fields import FluxConditioningField
+from invokeai.app.invocations.flux.flux_denoise import FluxDenoiseInvocation
+from invokeai.backend.flux.model import Flux, FluxParams
+from invokeai.backend.flux.schedulers import FLUX_SCHEDULER_MAP
+from invokeai.backend.model_manager.taxonomy import BaseModelType, FluxVariantType, ModelFormat, ModelType
+from invokeai.backend.patches.model_patch_raw import ModelPatchRaw
+from invokeai.backend.stable_diffusion.diffusion.conditioning_data import ConditioningFieldData, FLUXConditioningInfo
+from invokeai.backend.util.devices import TorchDevice
+
+
+class _ReachedTransformer(Exception):
+    pass
+
+
+class _CpuTimestepsEuler(FlowMatchEulerDiscreteScheduler):
+    def set_timesteps(self, *args, device=None, **kwargs):
+        # Meta has no scalar values for denoise's timestep.item(). Keep the real schedule on
+        # CPU while all session tensors still travel through the real meta concatenation.
+        return super().set_timesteps(*args, device="cpu", **kwargs)
+
+
+@pytest.mark.parametrize("route", ["fill", "control-lora"])
+@pytest.mark.parametrize("session_device", ["cpu", "meta"])
+def test_cpu_vae_conditioning_reaches_the_session_transformer(route, session_device, flux_shaped_vae, monkeypatch):
+    device = torch.device(session_device)
+    monkeypatch.setattr(TorchDevice, "choose_torch_device", lambda: device)
+    if session_device == "meta":
+        monkeypatch.setitem(FLUX_SCHEDULER_MAP, "euler", _CpuTimestepsEuler)
+    vae = flux_shaped_vae().eval()
+    encode = Mock(wraps=vae.encode)
+    monkeypatch.setattr(vae, "encode", encode)
+    vae_info = SimpleNamespace(
+        model=vae,
+        compute_device=torch.device("cpu"),
+        model_on_device=lambda **_: nullcontext((None, vae)),
+    )
+    # The real invocation and backend denoise run until the transformer forward. No CPU bf16
+    # matmul is needed: the forward is outside this VAE-to-session boundary regression.
+    transformer = Flux(
+        FluxParams(
+            in_channels=384 if route == "fill" else 128,
+            out_channels=64,
+            vec_in_dim=8,
+            context_in_dim=8,
+            hidden_size=12,
+            mlp_ratio=1.0,
+            num_heads=2,
+            depth=0,
+            depth_single_blocks=0,
+            axes_dim=[2, 2, 2],
+            theta=10000,
+            qkv_bias=True,
+            guidance_embed=True,
+        )
+    )
+
+    def transformer_input(**kwargs):
+        pixels = encode.call_args.args[0]
+        assert pixels.device == torch.device("cpu")
+        assert pixels.dtype == torch.float32
+        assert torch.isfinite(pixels).all()
+        img = kwargs["img"]
+        assert img.device == device
+        assert img.dtype == torch.bfloat16
+        assert img.shape == (1, 4, transformer.in_channels)
+        if route == "fill" and session_device == "cpu":
+            # Mask channels follow the 64 noise and 64 VAE channels, and survive packing.
+            assert (img[..., 128:] == 0).any()
+            assert (img[..., 128:] == 1).any()
+        raise _ReachedTransformer
+
+    monkeypatch.setattr(transformer, "forward", transformer_input)
+    transformer_info = SimpleNamespace(
+        model=transformer,
+        model_on_device=lambda **_: nullcontext((None, transformer)),
+    )
+    lora_info = SimpleNamespace(model=ModelPatchRaw(layers={}), model_in_ram=lambda: nullcontext())
+    models = {"vae": vae_info, "transformer": transformer_info, "lora": lora_info}
+    mask = torch.zeros(1, 32, 32, dtype=torch.bool)
+    mask[:, :, 16:] = True
+    context = SimpleNamespace(
+        models=SimpleNamespace(
+            load=models.__getitem__,
+            get_config=lambda _: SimpleNamespace(
+                base=BaseModelType.Flux,
+                type=ModelType.Main,
+                format=ModelFormat.Checkpoint,
+                variant=FluxVariantType.DevFill if route == "fill" else FluxVariantType.Dev,
+            ),
+        ),
+        images=SimpleNamespace(get_pil=lambda *_, **__: Image.new("RGB", (32, 32), (40, 80, 120))),
+        tensors=SimpleNamespace(load=lambda _: mask.clone()),
+        conditioning=SimpleNamespace(
+            load=lambda _: ConditioningFieldData(
+                conditionings=[FLUXConditioningInfo(clip_embeds=torch.zeros(1, 8), t5_embeds=torch.zeros(1, 2, 8))]
+            )
+        ),
+        logger=Mock(),
+    )
+    invocation = FluxDenoiseInvocation.model_construct(
+        transformer=SimpleNamespace(transformer="transformer", loras=[]),
+        controlnet_vae=SimpleNamespace(vae="vae"),
+        positive_text_conditioning=FluxConditioningField(conditioning_name="positive", mask=None),
+        control_lora=SimpleNamespace(img=SimpleNamespace(image_name="control"), lora="lora", weight=1.0)
+        if route == "control-lora"
+        else None,
+        fill_conditioning=SimpleNamespace(
+            image=SimpleNamespace(image_name="fill"), mask=SimpleNamespace(tensor_name="mask")
+        )
+        if route == "fill"
+        else None,
+        width=32,
+        height=32,
+        num_steps=1,
+        cfg_scale=1.0,
+        guidance=30.0,
+        seed=0,
+    )
+
+    with pytest.raises(_ReachedTransformer):
+        invocation._run_diffusion(context)
+
+    encode.assert_called_once()