Fix device issue

2026-03-11 05:29:48 +00:00 · 2025-06-26 19:14:25 -06:00
parent 4f91cb7148
commit 01a3c8a9b1
1 changed files with 0 additions and 1 deletions
--- a/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py
+++ b/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py
@@ -403,6 +403,5 @@ class FluxKontextModel(BaseModel):
                    
                control_latent = self.encode_images(control_tensor).to(latents.device, latents.dtype)
                latents = torch.cat((latents, control_latent), dim=1)
-                self.vae.to('cpu')

        return latents.detach()