From b2080da7c8a221c23b549fbecc1c649da49dfd7c Mon Sep 17 00:00:00 2001 From: aiXander Date: Fri, 19 Apr 2024 05:08:47 -0700 Subject: [PATCH] cleanup checkpoint --- trainer/checkpoint.py | 24 ++++++++++++++---------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/trainer/checkpoint.py b/trainer/checkpoint.py index 7f8de6e..fd6c706 100644 --- a/trainer/checkpoint.py +++ b/trainer/checkpoint.py @@ -95,19 +95,21 @@ def save_checkpoint( ) ) - # TODO this can prob be removed: - for idx, model in enumerate(text_encoder_peft_models): - if model is not None: - save_directory = os.path.join(output_dir, f"text_encoder_lora_{idx}") - model.save_pretrained( - save_directory = save_directory - ) - print(f"Saved text encoder {idx} in: {save_directory}") + # Separate txt_lora saves: + #for idx, model in enumerate(text_encoder_peft_models): + # if model is not None: + # save_directory = os.path.join(output_dir, f"text_encoder_lora_{idx}") + # model.save_pretrained( + # save_directory = save_directory + # ) + # print(f"Saved text encoder {idx} in: {save_directory}") if is_lora: assert len(unet_lora_parameters) > 0, f"Expected len(unet_lora_parameters) to be greater than zero if is_lora is True" - # TODO: I think we only need this to get the adapter_config.json? - unet.save_pretrained(save_directory = output_dir) + + # This saves adapter_config.json: + # TODO: adjust inference.py so it can load everything without needing this file + # unet.save_pretrained(save_directory = output_dir) text_encoder_lora_layers = [None, None] for idx, model in enumerate(text_encoder_peft_models): @@ -117,6 +119,8 @@ def save_checkpoint( lora_tensors = get_peft_model_state_dict(unet) unet_lora_layers_to_save = convert_state_dict_to_diffusers(lora_tensors) + + # TODO have a separate save for non SDXL model: StableDiffusionXLPipeline.save_lora_weights( output_dir, unet_lora_layers=unet_lora_layers_to_save,