only save 1 image per embedding (939f1652) · Commits · github_fork / Stable Diffusion Webui

modules/textual_inversion/textual_inversion.py

+4 −1

Original line number	Diff line number	Diff line
		@@ -242,6 +242,7 @@ def train_embedding(embedding_name, learn_rate, batch_size, data_root, log_direc

		last_saved_file = "<none>"
		last_saved_image = "<none>"
		embedding_yet_to_be_embedded = False

		ititial_step = embedding.step or 0
		if ititial_step > steps:
		@@ -281,6 +282,7 @@ def train_embedding(embedding_name, learn_rate, batch_size, data_root, log_direc
		if embedding.step > 0 and embedding_dir is not None and embedding.step % save_embedding_every == 0:
		last_saved_file = os.path.join(embedding_dir, f'{embedding_name}-{embedding.step}.pt')
		embedding.save(last_saved_file)
		embedding_yet_to_be_embedded = True

		write_loss(log_directory, "textual_inversion_loss.csv", embedding.step, len(ds), {
		"loss": f"{losses.mean():.7f}",
		@@ -318,7 +320,7 @@ def train_embedding(embedding_name, learn_rate, batch_size, data_root, log_direc

		shared.state.current_image = image

		if save_image_with_stored_embedding and os.path.exists(last_saved_file):
		if save_image_with_stored_embedding and os.path.exists(last_saved_file) and embedding_yet_to_be_embedded:

		last_saved_image_chunks = os.path.join(images_embeds_dir, f'{embedding_name}-{embedding.step}.png')

		@@ -342,6 +344,7 @@ def train_embedding(embedding_name, learn_rate, batch_size, data_root, log_direc
		captioned_image = insert_image_data_embed(captioned_image, data)

		captioned_image.save(last_saved_image_chunks, "PNG", pnginfo=info)
		embedding_yet_to_be_embedded = False

		image.save(last_saved_image)