github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Fix FLUX2 Klein load-time VRAM spikes on low-memory GPUs. (#726)

Browse Source 
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 <cursoragent@cursor.com>
Co-authored-by: Jaret Burkett <jaretburkett@gmail.com>
This commit is contained in:
M. Hofer
2026-04-01 17:36:55 +02:00
committed by GitHub
parent da2a79590f
commit f213e3b1e5
2 changed files with 14 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
extensions_built_in/diffusion_models/flux2/flux2_klein_model.py
View File

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

14
extensions_built_in/diffusion_models/flux2/flux2_model.py
View File

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