From cb046c28dbe5d179fc7b91b1dd2a5256ee7eae3e Mon Sep 17 00:00:00 2001 From: blepping Date: Tue, 21 Jan 2025 12:45:00 -0700 Subject: [PATCH] Fix Chroma guidance layer counting --- flux_mod/loader.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flux_mod/loader.py b/flux_mod/loader.py index 758120c..5ea9588 100644 --- a/flux_mod/loader.py +++ b/flux_mod/loader.py @@ -136,7 +136,7 @@ def load_flux_mod(model_path, timestep_guidance_path=None, linear_dtypes=torch.b keysplit[1].isnumeric() and keysplit[2] == "scale" ): - n_layers = int(keysplit[1]) + 1 + n_layers = max(n_layers, int(keysplit[1]) + 1) del state_dict[f"distilled_guidance_layer.{key}"] if n_layers == 0: raise RuntimeError("Could not determine number of distilled guidance layers in Chroma model")