From 8bfa50e2831f0f99fc8dc575cf1f3ff5e103e868 Mon Sep 17 00:00:00 2001 From: ddPn08 Date: Sun, 2 Apr 2023 00:10:19 +0900 Subject: [PATCH] small fix --- library/train_util.py | 6 ++++++ train_textual_inversion.py | 2 ++ train_textual_inversion_XTI.py | 2 ++ 3 files changed, 10 insertions(+) diff --git a/library/train_util.py b/library/train_util.py index 768d5b32..c6d49974 100644 --- a/library/train_util.py +++ b/library/train_util.py @@ -2851,6 +2851,8 @@ def save_sd_model_on_epoch_end( model_util.save_diffusers_checkpoint( args.v2, out_dir, text_encoder, unet, src_path, vae=vae, use_safetensors=use_safetensors ) + if args.huggingface_repo_id is not None: + huggingface_util.upload(out_dir, args, "/" + model_name) def remove_du(old_epoch_no): out_dir_old = os.path.join(args.output_dir, EPOCH_DIFFUSERS_DIR_NAME.format(model_name, old_epoch_no)) @@ -2906,6 +2908,8 @@ def save_sd_model_on_train_end( model_util.save_stable_diffusion_checkpoint( args.v2, ckpt_file, text_encoder, unet, src_path, epoch, global_step, save_dtype, vae ) + if args.huggingface_repo_id is not None: + huggingface_util.upload(ckpt_file, args, "/" + ckpt_name) else: out_dir = os.path.join(args.output_dir, model_name) os.makedirs(out_dir, exist_ok=True) @@ -2914,6 +2918,8 @@ def save_sd_model_on_train_end( model_util.save_diffusers_checkpoint( args.v2, out_dir, text_encoder, unet, src_path, vae=vae, use_safetensors=use_safetensors ) + if args.huggingface_repo_id is not None: + huggingface_util.upload(out_dir, args, "/" + model_name) def save_state_on_train_end(args: argparse.Namespace, accelerator): diff --git a/train_textual_inversion.py b/train_textual_inversion.py index c4b04554..3fb17f2e 100644 --- a/train_textual_inversion.py +++ b/train_textual_inversion.py @@ -493,6 +493,8 @@ def train(args): print(f"save trained model to {ckpt_file}") save_weights(ckpt_file, updated_embs, save_dtype) + if args.huggingface_repo_id is not None: + huggingface_util.upload(ckpt_file, args, "/" + ckpt_name) print("model saved.") diff --git a/train_textual_inversion_XTI.py b/train_textual_inversion_XTI.py index 58c79142..b0bc4c3a 100644 --- a/train_textual_inversion_XTI.py +++ b/train_textual_inversion_XTI.py @@ -537,6 +537,8 @@ def train(args): print(f"save trained model to {ckpt_file}") save_weights(ckpt_file, updated_embs, save_dtype) + if args.huggingface_repo_id is not None: + huggingface_util.upload(ckpt_file, args, "/" + ckpt_name) print("model saved.")