From b19a871aeb8622a791e5588223171bc896404dd7 Mon Sep 17 00:00:00 2001 From: linoytsaban Date: Fri, 22 Nov 2024 16:44:39 +0200 Subject: [PATCH 1/3] smol change to fix checkpoint saving & resuming (as done in train_dreambooth_sd3.py) --- examples/dreambooth/train_dreambooth_lora_sd3.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/examples/dreambooth/train_dreambooth_lora_sd3.py b/examples/dreambooth/train_dreambooth_lora_sd3.py index dcf093a94c5a..933ec122258f 100644 --- a/examples/dreambooth/train_dreambooth_lora_sd3.py +++ b/examples/dreambooth/train_dreambooth_lora_sd3.py @@ -1294,10 +1294,13 @@ def save_model_hook(models, weights, output_dir): for model in models: if isinstance(model, type(unwrap_model(transformer))): transformer_lora_layers_to_save = get_peft_model_state_dict(model) - elif isinstance(model, type(unwrap_model(text_encoder_one))): - text_encoder_one_lora_layers_to_save = get_peft_model_state_dict(model) - elif isinstance(model, type(unwrap_model(text_encoder_two))): - text_encoder_two_lora_layers_to_save = get_peft_model_state_dict(model) + elif isinstance(model, type(unwrap_model(text_encoder_one))): # or text_encoder_two + # check hidden size to distinguish between text_encoder_one and two + hidden_size = unwrap_model(model).config.hidden_size + if hidden_size == 768: + text_encoder_one_lora_layers_to_save = get_peft_model_state_dict(model) + elif hidden_size == 1280: + text_encoder_two_lora_layers_to_save = get_peft_model_state_dict(model) else: raise ValueError(f"unexpected save model: {model.__class__}") From 6c1cef1a88e86af2217f9ff80ef0ace82d24e97a Mon Sep 17 00:00:00 2001 From: Linoy Date: Fri, 22 Nov 2024 14:52:24 +0000 Subject: [PATCH 2/3] style --- examples/dreambooth/train_dreambooth_lora_sd3.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/dreambooth/train_dreambooth_lora_sd3.py b/examples/dreambooth/train_dreambooth_lora_sd3.py index 933ec122258f..acdfbfac7c58 100644 --- a/examples/dreambooth/train_dreambooth_lora_sd3.py +++ b/examples/dreambooth/train_dreambooth_lora_sd3.py @@ -1294,7 +1294,7 @@ def save_model_hook(models, weights, output_dir): for model in models: if isinstance(model, type(unwrap_model(transformer))): transformer_lora_layers_to_save = get_peft_model_state_dict(model) - elif isinstance(model, type(unwrap_model(text_encoder_one))): # or text_encoder_two + elif isinstance(model, type(unwrap_model(text_encoder_one))): # or text_encoder_two # check hidden size to distinguish between text_encoder_one and two hidden_size = unwrap_model(model).config.hidden_size if hidden_size == 768: From 613a32d0e6a23f337a705a4a68833e02d3be9c50 Mon Sep 17 00:00:00 2001 From: linoytsaban Date: Sun, 24 Nov 2024 17:14:52 +0200 Subject: [PATCH 3/3] modify comment to explain reasoning behind hidden size check --- examples/dreambooth/train_dreambooth_lora_sd3.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/dreambooth/train_dreambooth_lora_sd3.py b/examples/dreambooth/train_dreambooth_lora_sd3.py index acdfbfac7c58..3f721e56addf 100644 --- a/examples/dreambooth/train_dreambooth_lora_sd3.py +++ b/examples/dreambooth/train_dreambooth_lora_sd3.py @@ -1295,7 +1295,7 @@ def save_model_hook(models, weights, output_dir): if isinstance(model, type(unwrap_model(transformer))): transformer_lora_layers_to_save = get_peft_model_state_dict(model) elif isinstance(model, type(unwrap_model(text_encoder_one))): # or text_encoder_two - # check hidden size to distinguish between text_encoder_one and two + # both text encoders are of the same class, so we check hidden size to distinguish between the two hidden_size = unwrap_model(model).config.hidden_size if hidden_size == 768: text_encoder_one_lora_layers_to_save = get_peft_model_state_dict(model)