VictorMylle/Kohya-ss-sd-scripts
1
0
Fork 1
mirror of https://github.com/kohya-ss/sd-scripts.git synced 2026-04-08 22:35:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

Added feature to upload to huggingface

Browse Source
This commit is contained in:
ddPn08
2023-03-28 00:49:09 +09:00
parent c639cb7d5d
commit d42431d73a
3 changed files with 78 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
train_network.py
View File

@@ -24,6 +24,7 @@ from library.config_util import (
ConfigSanitizer,
BlueprintGenerator,
)
import library.utils as utils
import library.custom_train_functions as custom_train_functions
from library.custom_train_functions import apply_snr_weight
@@ -626,6 +627,7 @@ def train(args):
metadata["ss_training_finished_at"] = str(time.time())
print(f"saving checkpoint: {ckpt_file}")
unwrap_model(network).save_weights(ckpt_file, save_dtype, minimum_metadata if args.no_metadata else metadata)
utils.huggingface_upload(ckpt_file, args, "/" + ckpt_name)
def remove_old_func(old_epoch_no):
old_ckpt_name = train_util.EPOCH_FILE_NAME.format(model_name, old_epoch_no) + "." + args.save_model_as
@@ -665,6 +667,7 @@ def train(args):
print(f"save trained model to {ckpt_file}")
network.save_weights(ckpt_file, save_dtype, minimum_metadata if args.no_metadata else metadata)
utils.huggingface_upload(ckpt_file, args, "/" + ckpt_name)
print("model saved.")