Fix: compute device partial load

[Pfannkuchensack]

Summary

fix: VAE decode — and, in the same class of bug, text encoding — could silently run entirely on the CPU when partial loading is active, even on machines with a configured CUDA/MPS device.

Why: Under partial loading with VRAM pressure, the model cache can temporarily offload all of a model's weights back to RAM. The affected invocations inferred their execution device from the model's current parameter residency via get_effective_device(model), which returns cpu in that state. The latents/tokens were then placed on the CPU, and because the custom autocast layers cast their weights to the input's device, the whole forward pass ran on the CPU — very slow, and for the SD/SDXL VAE it additionally forced fp32. The bug is environment-dependent (only triggers when no weight happens to stay resident on the accelerator), which is why it reproduces inconsistently.

This was introduced by #9293 (VAE) — commit 82e2681126 — which replaced TorchDevice.choose_torch_device() with get_effective_device(vae). The text-encoder invocations (from #8777) shared the same trap for models whose modules are all autocast-capable (e.g. CLIP), where repair_required_tensors_on_device() pins nothing to the compute device.

How: Use the model's intended compute device instead of inferring it from residency:

• Added LoadedModel.compute_device — stable regardless of partial-load state, and still honors cpu_only.
• VAE decode (8 invocations) and text encoders (compel/SD3/z_image/flux2_klein/cogview4/qwen_image) now use *_info.compute_device.
• HFEncoder (FLUX conditioner) gained an optional device param (fallback = old behavior); the FLUX text-encoder invocation threads compute_device through for T5 & CLIP.
• The repair_required_tensors_on_device() calls are unchanged — they solve a separate problem.

Related Issues / Discussions

Closes #9373

The issue only reports the VAE path; the text-encoder fix is the same root cause and is validated below.

QA Instructions

Repro (before this PR):

1. Enable partial loading; set device_working_mem_gb to 4.0.
2. Use an SDXL model + VAE; clear the model cache and generate under enough VRAM pressure that the VAE gets fully offloaded.
3. Decode at large resolution with FP32 VAE precision.
4. Observe the VAE decoder running on the CPU (slow) instead of the GPU.

With this PR: the decode/encode runs on the configured compute device regardless of partial-load state; cpu_only-configured VAEs/encoders still correctly run on the CPU.

Automated: new regression test tests/backend/model_manager/load/test_loaded_model_compute_device.py reproduces the offloaded-weights case for both a VAE-like conv model and a fully-autocast (CLIP-like) encoder. Updated test_compel, test_sd3_text_encoder, test_cogview4_text_encoder, test_conditioner to assert the intended-device behavior. All green (10 passed, 2 skipped [mps]); ruff clean.

Merge Plan

Standard merge — no DB/schema changes, no migrations.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Tests added / updated (if applicable)
• ❗Changes to a redux slice have a corresponding migration — n/a (no redux changes)
• Documentation added / updated (if applicable) — n/a
• Updated What's New copy (if doing a release after this PR)

[DustyShoe]

That did fix the issue, Thank you!

[lstein]

Approving — well-scoped fix for a real, environment-dependent bug

The root-cause analysis is correct: under partial loading with VRAM pressure, all of a model's weights can be temporarily offloaded to RAM, so get_effective_device(model) (and HF model.device) report cpu. Because the custom autocast layers follow the input device, the whole decode/encode then silently runs on the CPU — slow, and for the SD/SDXL VAE it also forces fp32. Replacing residency inference with the stable, intent-based LoadedModel.compute_device is the right fix.

What I verified:

• LoadedModel.compute_device cleanly delegates to cached_model.compute_device, which exists on both CachedModelOnlyFullLoad and CachedModelWithPartialLoad and correctly honors cpu_only (checked against model_cache.put's execution_device plumbing). No behavioral change in the common fully-resident case.
• HFEncoder's new device param is backward-compatible (defaults to the old get_effective_device fallback).
• All 8 VAE decodes + the text encoders (compel/SD3/z_image/flux2_klein/cogview4/qwen_image/FLUX) are converted consistently; the SD3 refactor to hoist t5_text_encoder_info is correct.
• The tests are excellent — the torch.device("meta") stand-in for an accelerator lets the offloaded-residency mechanism be reproduced on CPU-only CI, and both #9373 variants (conv VAE + fully-autocast CLIP-like) are exercised end-to-end. Ran the affected suite locally: green, ruff clean.

One gap I found and fixed (pushed to this branch)

anima_text_encoder.py still inferred its device from residency:

device = text_encoder.device   # HF PreTrainedModel residency → CPU under full offload

This is the same class of bug the PR sets out to eliminate. The Qwen3 0.6B encoder is fully autocast-capable and this path doesn't call repair_required_tensors_on_device(), so under a full offload text_encoder.device returns cpu and the whole encode runs there. The impact is slightly broader here because device also feeds TorchDevice.choose_anima_inference_dtype(device). Notably the PR already fixed the anima VAE decode but left the anima text encoder on the buggy pattern.

Fix (commit ed77554f79):

device = text_encoder_info.compute_device

plus a new tests/app/invocations/test_anima_text_encoder.py mirroring the SD3/compel tests — covering both the offloaded-accelerator case and the cpu_only case. I confirmed the test actually pins the behavior by temporarily reverting the fix (it fails) before restoring it.

With that, the "same class of bug for text encoders" claim is complete across every encoder. Nice catch on the root cause. 🚀