[wunianze666-netizen]

Summary

VAE encoding inputs, memory estimates and tile decisions follow the loaded model's compute_device. Encoded conditioning is then transferred to the device and dtype of its consuming transformer.

The FLUX Fill and Control LoRA paths also handle a CPU VAE in an accelerator session. Fill keeps its pixels and boolean mask together before encoding, aligns the packed mask with the returned latents, and transfers the assembled conditioning to the session. Control LoRA transfers its VAE result before denoise concatenation.

Related Issues / Discussions

Closes #10014. Also covers the two additional FLUX routes identified in the review.

QA Instructions

The finalized follow-up fixtures fail on unchanged 23cd5f1 in four cases: Fill and Control LoRA, each under CPU/meta session conditions. The two CPU cases expose a dtype mismatch; the meta cases expose an early pixel transfer or a cross-device concatenation.

On this revision, 11 focused cases and 113 related cases pass across disjoint files (124 unique passes). The new four-case FLUX file also passes independently; those four are included in the 11. A real tiny fp32 VAE encodes with its unused decoder offloaded to meta, so using the first parameter's device fails inside the native encoder. Separate deliberate regression controls fail when that incorrect device lookup or any of the three estimator device arguments is substituted.

Validation used Windows, Python 3.12.14, Torch 2.13.0+cpu and the project's pinned Diffusers 0.42.1 installed in an isolated import target. Native pytest ran with -n 0, bounded BLAS threads and isolated test roots. This is not the complete uv-locked environment. Ruff 0.11.2 lint/format checks, git diff --check and normal commit hooks pass.

VAE forwards use CPU fp32. Meta sessions and the transformer input boundary exercise routing without accelerator kernels. No physical GPU, full-model generation, real partial-cache loading, family E2E, RSS or speed result is claimed.

Review

The device contract uses LoadedModel.compute_device, independently of current weight residency. Estimator spies verify both the device and reservation forwarding through the existing fixed_working_memory fixture.

Fill defers dtype conversion to the VAE owner. For fp32/fp16 VAEs this removes the former premature bf16 quantization, so those inputs are not claimed numerically identical to the old path. Boolean mask semantics remain unchanged.

The shared is_oom_error classifier recognizes the two known PyTorch CPU allocator message formats. That classification applies to its eleven existing consumers, including Kontext's tile retry; no unrelated handler or retry policy is changed.

Implementation and validation were assisted by OpenAI Codex.

Checklist

• The PR has a short but descriptive title, suitable for a changelog
• Meaningful regression coverage added / updated where needed; obsolete tests/code removed
• Persisted-state and API changes include required migrations / compatibility validation
• Relevant performance/efficiency opportunities considered; material claims have evidence
• Material review findings resolved and relevant checks rerun
• Documentation added / updated (if applicable)
• Updated What's New copy (if doing a release after this PR)

---

^{Open in [code]smith · Need help on this PR? Tag @codesmith-bot with what you need.}

[Pfannkuchensack]

Thanks, this fixes every path listed in #10014. Encode, tile decision and estimate all follow compute_device, and results are moved on to where downstream expects them. Reverting any listed path makes a test fail, also on a CPU-only box. No blockers; checked on CPU with a simulated accelerator session, not on real hardware.

Should fix

1. CPU bf16 in a test. tests/app/invocations/test_vae_compute_device.py:46 runs the CogView4 case in bf16 on CPU. tests/AGENTS.md rules out CPU bf16 matmul because it can abort the worker on part of the Windows fleet. These are convolutions and windows-cpu passed, but low precision isn't what's under test here. Please run it in fp32 and patch the estimator with the existing fixed_working_memory fixture.
2. compute_device isn't told apart from weight residency. Every test has compute_device=cpu with CPU weights, so replacing the fix with next(vae.parameters()).device passes on sd3, ernie, anima, kontext and flux2. That is the exact wrong approach the issue warns about for partial loading. One case with compute_device different from the weights would close it.

Follow-up, not caused by this PR. A cpu_only FLUX VAE still crashes on two routes the issue didn't list: _prep_flux_fill_img_cond (flux_denoise.py:908-919) and the Control LoRA conditioning (:856-859, then flux/denoise.py:151). Both use the latents from vae_encode without moving them. Two .to(device, dtype) calls here, or a follow-up issue, would both work.

Nits

• is_oom_error in oom.py:35 is shared by eleven call sites, while the description names only Kontext's retry.
• The QA section cites Diffusers 0.41.0; the pin is 0.42.1.
• Dropping device=vae_info.compute_device from the sd3, cogview4 and qwen estimators passes every test.

[wunianze666-netizen]

Addressed the review in 857fb06.

• CogView4 now uses a real fp32 CPU VAE with fixed_working_memory. The residency regression offloads the unused decoder to meta while encoding on CPU; separate controls fail for the first-parameter device lookup and each omitted estimator device argument.
• Fixed both additional FLUX routes. Fill needed the early session-device pixel transfer removed as well as the final conditioning transfer; otherwise a CPU VAE still received meta/session pixels. The mask follows the pixels and returned latents. Control LoRA transfers its VAE result before concatenation.
• Final fixtures give four failures on unchanged 23cd5f1 and pass here. The focused 11 and related 113 cases are disjoint: 124 unique passes, with the standalone four-case file included in that total. Ruff 0.11.2 and normal hooks pass.

The QA now records isolated Diffusers 0.42.1, CPU/meta limits, the Fill dtype-conversion change, and the shared OOM classifier's scope across eleven consumers. No physical accelerator or full-model result is claimed.