| Fix: compute device partial load (#9375) * fix(vae): decode on the model's intended compute device, not current residency (#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 #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 #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 #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> | 2 个月前 |