From 01a3c8a9b1946695884f6df00e7037c6a748b25c Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Thu, 26 Jun 2025 19:14:25 -0600 Subject: [PATCH] Fix device issue --- .../diffusion_models/flux_kontext/flux_kontext.py | 1 - 1 file changed, 1 deletion(-) diff --git a/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py b/extensions_built_in/diffusion_models/flux_kontext/flux_kontext.py index da6acb94..8d4247c9 100644 --- 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() \ No newline at end of file