Add sample batch size for Lumina

This commit is contained in:
rockerBOO
2025-02-23 20:19:24 -05:00
parent ba725a84e9
commit 48e7da2d4a
3 changed files with 201 additions and 106 deletions

View File

@@ -1242,6 +1242,7 @@ class NetworkTrainer:
# For --sample_at_first
optimizer_eval_fn()
self.sample_images(accelerator, args, 0, global_step, accelerator.device, vae, tokenizers, text_encoder, unet)
progress_bar.unpause() # Reset progress bar to before sampling images
optimizer_train_fn()
is_tracking = len(accelerator.trackers) > 0
if is_tracking:
@@ -1344,6 +1345,7 @@ class NetworkTrainer:
self.sample_images(
accelerator, args, None, global_step, accelerator.device, vae, tokenizers, text_encoder, unet
)
progress_bar.unpause()
# 指定ステップごとにモデルを保存
if args.save_every_n_steps is not None and global_step % args.save_every_n_steps == 0:
@@ -1531,6 +1533,7 @@ class NetworkTrainer:
train_util.save_and_remove_state_on_epoch_end(args, accelerator, epoch + 1)
self.sample_images(accelerator, args, epoch + 1, global_step, accelerator.device, vae, tokenizers, text_encoder, unet)
progress_bar.unpause()
optimizer_train_fn()
# end of epoch