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

Update train_network.py

Browse Source
This commit is contained in:
DKnight54
2025-02-03 16:37:44 +08:00
committed by GitHub
parent 4c5ec694f7
commit e6433ec035
1 changed files with 5 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
train_network.py
View File

@@ -132,8 +132,8 @@ class NetworkTrainer:
if param.grad is not None:
param.grad = accelerator.reduce(param.grad, reduction="mean")
def sample_images(self, accelerator, args, epoch, global_step, device, vae, tokenizer, text_encoder, unet, latents_list=None):
train_util.sample_images(accelerator, args, epoch, global_step, device, vae, tokenizer, text_encoder, unet, latents_list)
def sample_images(self, accelerator, args, epoch, global_step, device, vae, tokenizer, text_encoder, unet, example_tuple=None):
train_util.sample_images(accelerator, args, epoch, global_step, device, vae, tokenizer, text_encoder, unet, example_tuple)
def train(self, args):
session_id = random.randint(0, 2**32)
@@ -1023,24 +1023,13 @@ class NetworkTrainer:
keys_scaled, mean_norm, maximum_norm = None, None, None
# Checks if the accelerator has performed an optimization step behind the scenes
# Collecting latents and caption lists from all processes
example_tuple = (latents, batch["captions"])
# Converts batch of latents into list of dicts containing individual latents, height and width to merge across processes
# Allows for different latent sizes
latents_list = []
for idx in range(len(batch["captions"])):
latent_dict = {}
latent_dict["prompt"] = batch["captions"][idx]
latent_dict["height"] = latents.shape[2] * 8
latent_dict["width"] = latents.shape[3] * 8
latent_dict["original_lantent"] = latents[idx].unsqueeze(0)
latents_list.append(latent_dict)
latents_list = gather_object(latents_list)
if accelerator.sync_gradients:
progress_bar.update(1)
global_step += 1
self.sample_images(accelerator, args, None, global_step, accelerator.device, vae, tokenizer, text_encoder, unet, latents_list)
self.sample_images(accelerator, args, None, global_step, accelerator.device, vae, tokenizer, text_encoder, unet, example_tuple)
# 指定ステップごとにモデルを保存
if args.save_every_n_steps is not None and global_step % args.save_every_n_steps == 0:
@@ -1096,7 +1085,7 @@ class NetworkTrainer:
if args.save_state:
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, tokenizer, text_encoder, unet, latents_list)
self.sample_images(accelerator, args, epoch + 1, global_step, accelerator.device, vae, tokenizer, text_encoder, unet, example_tuple)
# end of epoch