From 7602e476ebf241ead1a9b5b43290a3ee5baf3d67 Mon Sep 17 00:00:00 2001 From: Jaret Burkett Date: Mon, 13 Jul 2026 10:14:53 -0600 Subject: [PATCH] Exclude sensative layers from quantization in krea --- .../diffusion_models/krea2/krea2.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/extensions_built_in/diffusion_models/krea2/krea2.py b/extensions_built_in/diffusion_models/krea2/krea2.py index 0fa8cc9..128fdb0 100644 --- a/extensions_built_in/diffusion_models/krea2/krea2.py +++ b/extensions_built_in/diffusion_models/krea2/krea2.py @@ -379,6 +379,24 @@ class Krea2Model(BaseModel): # tell the model to invert assistant on inference since we want remove lora effects self.invert_assistant_lora = True + + def get_quantization_exclude_modules(self): + # sensitive modules kept in full precision (fnmatch patterns on module + # names within SingleStreamDiT): + # first - patchified latent input projection + # tmlp* / tproj* - timestep embedder + modulation projection; feed + # every block's DoubleSharedModulation and LastLayer + # txtmlp* - text feature -> model width projection + # txtfusion.projector - tiny (num_txt_layers -> 1) encoder-layer mixer + # last* - final norm/modulated output projection + return [ + "first", + "tmlp*", + "tproj*", + "txtmlp*", + "txtfusion.projector", + "last*", + ] def load_model(self): dtype = self.torch_dtype