Studio: build VAE tile blend weights on CPU so MPS tiled encode works - #12937

Confirmed main still builds the float64 blend weights on the VAE device in diffusion_vae_tiling.py, which MPS rejects on every tiled encode, and moving only the fp32 result over looks right. Will get this reviewed, ideally with a run on a real Apple Silicon machine.

@gokay-ai Thank you for the PR. I tested this on Apple Silicon and verified it works.

Neither red check is from this PR. Repo tests (CPU, studio) fails test_embedding_model_controls_stack_on_the_narrowest_viewports, red on main since #12875 and fixed in #12945. GGUF inference smoke stalled in the tool-calling step after the bash tool passed (no output from 11:49 until the 20 minute timeout); I have re-run it.

Pushed two commits on top of the original fix, plus a merge of the latest main.

What changed

• axis_weights now decides where to build its float64 values from the device's supports_float64 capability (new float64_device helper in diffusion_device.py), the same flag force_float32_rope already keys off. A device with float64 builds in place exactly as main does, so CUDA / XPU / CPU run the same operations as before; only MPS builds on CPU and receives the fp32 result.
• The LTX-2 axis_weights was a copy of the image one that differed only in its default margin and ramp. It now delegates to the image implementation, so the float64 handling lives in one place.
• The regression test records the dtype and device of every tensor the build produces instead of the constructor keyword arguments. The earlier spy test still passed with .to(device).float() in place of .float().to(device), which puts float64 on the device.

Validation

• Apple Silicon (MPS, torch 2.10): on main, axis_weights, the tiled encode and the tiled decode of a small random AutoencoderKLQwenImage21 all raise the float64 TypeError; with this branch all three run, and the weights are fp32 on mps and bit-equal to the CPU weights.
• CUDA (RTX 4070 SUPER, torch 2.11): weights built in float64 on cuda are bit-equal to weights built on CPU across 1592 tile geometries (image decode 8x / 16x / 32x, image encode, LTX-2), so moving the LTX-2 build onto the device changes no output.

cd studio/backend
python -m pytest tests/test_diffusion_vae_tiling.py tests/test_video_ltx2_vae_tiles.py tests/test_diffusion_device.py tests/test_vae_tile_geometry.py -q
# 816 passed

Review and CI status for head e9ca0e2d4b, run on a staging PR with the same head.

• Codex review: "Didn't find any major issues" on e9ca0e2d4b, first round, no inline findings: [staging] #12937 Studio: build VAE tile blend weights on CPU so MPS tiled encode works Lyxot/unsloth#145 (comment)
• CI: 38 passed, 2 skipped, 1 failed: https://github.com/Lyxot/unsloth/pull/145/checks
• The one failure is not from this PR. (Python 3.13, a-k) fails seven tests in tests/test_hub_token_caller_identity.py (TypeError: _st_resolver.<locals>.<lambda>() takes 1 positional argument but 2 were given). This branch does not touch that file, the same seven fail on main at d8ad0b0806, and Hub token identity tests: let the llama-backend stubs take the caller's token after #13005 #13022 is open to fix them.
• Changed during this pass: nothing beyond what the previous comment describes (the two commits and the merge of main). There were no unresolved review threads, and no further commits were needed.