From affa411edc916b55d140812730a3bfc4d4d2d2e9 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Sun, 29 Mar 2026 09:49:10 -0600 Subject: [PATCH] Fixed an issue where Flux.2 model VAE can be left offloaded to CPU when encoding control images while caching latents --- extensions_built_in/diffusion_models/flux2/flux2_model.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/extensions_built_in/diffusion_models/flux2/flux2_model.py b/extensions_built_in/diffusion_models/flux2/flux2_model.py index 1e77474..c07cde4 100644 --- a/extensions_built_in/diffusion_models/flux2/flux2_model.py +++ b/extensions_built_in/diffusion_models/flux2/flux2_model.py @@ -391,6 +391,8 @@ class Flux2Model(BaseModel): control_img = control_img * 2 - 1 controls.append(control_img) + if self.vae.device == torch.device("cpu"): + self.vae.to(self.device_torch) img_cond_seq_item, img_cond_seq_ids_item = encode_image_refs( self.vae, controls, limit_pixels=control_image_max_res )