From f213e3b1e5ab591543b49222916cbabbbafff7fd Mon Sep 17 00:00:00 2001 From: "M. Hofer" <20087022+mhofer1976@users.noreply.github.com> Date: Wed, 1 Apr 2026 17:36:55 +0200 Subject: [PATCH] Fix FLUX2 Klein load-time VRAM spikes on low-memory GPUs. (#726) Keep the transformer and Qwen text encoder off CUDA during initial load/quantization in low-VRAM mode so model startup avoids full-model OOM before offloading and quantization can take effect. Co-authored-by: Cursor Co-authored-by: Jaret Burkett --- .../diffusion_models/flux2/flux2_klein_model.py | 9 ++++----- .../diffusion_models/flux2/flux2_model.py | 14 ++++++++++---- 2 files changed, 14 insertions(+), 9 deletions(-) diff --git a/extensions_built_in/diffusion_models/flux2/flux2_klein_model.py b/extensions_built_in/diffusion_models/flux2/flux2_klein_model.py index 86fdafa..864bb41 100644 --- a/extensions_built_in/diffusion_models/flux2/flux2_klein_model.py +++ b/extensions_built_in/diffusion_models/flux2/flux2_klein_model.py @@ -44,15 +44,14 @@ class Flux2KleinModel(Flux2Model): self.flux2_klein_te_path, torch_dtype=dtype, ) - text_encoder.to(self.device_torch, dtype=dtype) - - flush() - if self.model_config.quantize_te: self.print_and_status_update("Quantizing Qwen3") - quantize(text_encoder, weights=get_qtype(self.model_config.qtype)) + quantize(text_encoder, weights=get_qtype(self.model_config.qtype_te)) freeze(text_encoder) flush() + elif not self.model_config.low_vram: + text_encoder.to(self.device_torch, dtype=dtype) + flush() if ( self.model_config.layer_offloading diff --git a/extensions_built_in/diffusion_models/flux2/flux2_model.py b/extensions_built_in/diffusion_models/flux2/flux2_model.py index c07cde4..2a09f81 100644 --- a/extensions_built_in/diffusion_models/flux2/flux2_model.py +++ b/extensions_built_in/diffusion_models/flux2/flux2_model.py @@ -155,11 +155,11 @@ class Flux2Model(BaseModel): transformer.load_state_dict(transformer_state_dict, assign=True) - transformer.to(self.quantize_device, dtype=dtype) - if self.model_config.quantize: # patch the state dict method patch_dequantization_on_save(transformer) + # Avoid full-model peak VRAM allocation before quantization. + self.print_and_status_update("Keeping transformer on CPU for quantization") self.print_and_status_update("Quantizing Transformer") quantize_model(self, transformer) flush() @@ -234,10 +234,16 @@ class Flux2Model(BaseModel): flush() # just to make sure everything is on the right device and dtype - text_encoder[0].to(self.device_torch) + if self.model_config.low_vram: + text_encoder[0].to("cpu") + else: + text_encoder[0].to(self.device_torch) text_encoder[0].requires_grad_(False) text_encoder[0].eval() - pipe.transformer = pipe.transformer.to(self.device_torch) + if self.model_config.low_vram: + pipe.transformer = pipe.transformer.to("cpu") + else: + pipe.transformer = pipe.transformer.to(self.device_torch) flush() # save it to the model class