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

@@ -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
                     )