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

add callback for step start

Browse Source
This commit is contained in:
Kohya S
2023-05-11 22:00:41 +09:00
parent 8d562ecf48
commit 7889a52f95
1 changed files with 6 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
train_network.py
View File

@@ -525,10 +525,11 @@ def train(args):
loss_total = 0.0
del train_dataset_group
# if hasattr(network, "on_step_start"):
# on_step_start = network.on_step_start
# else:
# on_step_start = lambda *args, **kwargs: None
# callback for step start
if hasattr(network, "on_step_start"):
on_step_start = network.on_step_start
else:
on_step_start = lambda *args, **kwargs: None
# function for saving/removing
def save_model(ckpt_name, unwrapped_nw, steps, epoch_no, force_sync_upload=False):
@@ -563,7 +564,7 @@ def train(args):
for step, batch in enumerate(train_dataloader):
current_step.value = global_step
with accelerator.accumulate(network):
# on_step_start(text_encoder, unet)
on_step_start(text_encoder, unet)
with torch.no_grad():
if "latents" in batch and batch["latents"] is not None: