Skip to content

S
Stable Diffusion Webui
Commit 3ce2bfdf authored by Muhammad Rizqi Nur's avatar Muhammad Rizqi Nur
Browse files
Options

Add cleanup after training

parent ab27c111
Hide whitespace changes
Inline Side-by-side
Showing with 200 additions and 186 deletions
modules/hypernetworks/hypernetwork.py
View file @ 3ce2bfdf
...@@ -398,110 +398,112 @@ def train_hypernetwork(hypernetwork_name, learn_rate, batch_size, data_root, log ...@@ -398,110 +398,112 @@ def train_hypernetwork(hypernetwork_name, learn_rate, batch_size, data_root, log
forced_filename = "<none>" forced_filename = "<none>"
pbar = tqdm.tqdm(enumerate(ds), total=steps - ititial_step) pbar = tqdm.tqdm(enumerate(ds), total=steps - ititial_step)
for i, entries in pbar:
hypernetwork.step = i + ititial_step
if len(loss_dict) > 0:
previous_mean_losses = [i[-1] for i in loss_dict.values()]
previous_mean_loss = mean(previous_mean_losses)
scheduler.apply(optimizer, hypernetwork.step)
if scheduler.finished:
break
if shared.state.interrupted:
break
with torch.autocast("cuda"):
c = stack_conds([entry.cond for entry in entries]).to(devices.device)
# c = torch.vstack([entry.cond for entry in entries]).to(devices.device)
x = torch.stack([entry.latent for entry in entries]).to(devices.device)
loss = shared.sd_model(x, c)[0]
del x
del c
losses[hypernetwork.step % losses.shape[0]] = loss.item()
for entry in entries:
loss_dict[entry.filename].append(loss.item())
optimizer.zero_grad()
weights[0].grad = None
loss.backward()
if weights[0].grad is None: try:
steps_without_grad += 1 for i, entries in pbar:
hypernetwork.step = i + ititial_step
if len(loss_dict) > 0:
previous_mean_losses = [i[-1] for i in loss_dict.values()]
previous_mean_loss = mean(previous_mean_losses)
scheduler.apply(optimizer, hypernetwork.step)
if scheduler.finished:
break
if shared.state.interrupted:
break
with torch.autocast("cuda"):
c = stack_conds([entry.cond for entry in entries]).to(devices.device)
# c = torch.vstack([entry.cond for entry in entries]).to(devices.device)
x = torch.stack([entry.latent for entry in entries]).to(devices.device)
loss = shared.sd_model(x, c)[0]
del x
del c
losses[hypernetwork.step % losses.shape[0]] = loss.item()
for entry in entries:
loss_dict[entry.filename].append(loss.item())
optimizer.zero_grad()
weights[0].grad = None
loss.backward()
if weights[0].grad is None:
steps_without_grad += 1
else:
steps_without_grad = 0
assert steps_without_grad < 10, 'no gradient found for the trained weight after backward() for 10 steps in a row; this is a bug; training cannot continue'
optimizer.step()
steps_done = hypernetwork.step + 1
if torch.isnan(losses[hypernetwork.step % losses.shape[0]]):
raise RuntimeError("Loss diverged.")
if len(previous_mean_losses) > 1:
std = stdev(previous_mean_losses)
else: else:
steps_without_grad = 0 std = 0
assert steps_without_grad < 10, 'no gradient found for the trained weight after backward() for 10 steps in a row; this is a bug; training cannot continue' dataset_loss_info = f"dataset loss:{mean(previous_mean_losses):.3f}" + u"\u00B1" + f"({std / (len(previous_mean_losses) ** 0.5):.3f})"
pbar.set_description(dataset_loss_info)
optimizer.step()
if hypernetwork_dir is not None and steps_done % save_hypernetwork_every == 0:
steps_done = hypernetwork.step + 1 # Before saving, change name to match current checkpoint.
hypernetwork.name = f'{hypernetwork_name}-{steps_done}'
if torch.isnan(losses[hypernetwork.step % losses.shape[0]]): last_saved_file = os.path.join(hypernetwork_dir, f'{hypernetwork.name}.pt')
raise RuntimeError("Loss diverged.") hypernetwork.save(last_saved_file)
if len(previous_mean_losses) > 1: textual_inversion.write_loss(log_directory, "hypernetwork_loss.csv", hypernetwork.step, len(ds), {
std = stdev(previous_mean_losses) "loss": f"{previous_mean_loss:.7f}",
else: "learn_rate": scheduler.learn_rate
std = 0 })
dataset_loss_info = f"dataset loss:{mean(previous_mean_losses):.3f}" + u"\u00B1" + f"({std / (len(previous_mean_losses) ** 0.5):.3f})"
pbar.set_description(dataset_loss_info) if images_dir is not None and steps_done % create_image_every == 0:
forced_filename = f'{hypernetwork_name}-{steps_done}'
if hypernetwork_dir is not None and steps_done % save_hypernetwork_every == 0: last_saved_image = os.path.join(images_dir, forced_filename)
# Before saving, change name to match current checkpoint.
hypernetwork.name = f'{hypernetwork_name}-{steps_done}' optimizer.zero_grad()
last_saved_file = os.path.join(hypernetwork_dir, f'{hypernetwork.name}.pt') shared.sd_model.cond_stage_model.to(devices.device)
hypernetwork.save(last_saved_file) shared.sd_model.first_stage_model.to(devices.device)
textual_inversion.write_loss(log_directory, "hypernetwork_loss.csv", hypernetwork.step, len(ds), { p = processing.StableDiffusionProcessingTxt2Img(
"loss": f"{previous_mean_loss:.7f}", sd_model=shared.sd_model,
"learn_rate": scheduler.learn_rate do_not_save_grid=True,
}) do_not_save_samples=True,
)
if images_dir is not None and steps_done % create_image_every == 0:
forced_filename = f'{hypernetwork_name}-{steps_done}'
last_saved_image = os.path.join(images_dir, forced_filename)
optimizer.zero_grad()
shared.sd_model.cond_stage_model.to(devices.device)
shared.sd_model.first_stage_model.to(devices.device)
p = processing.StableDiffusionProcessingTxt2Img(
sd_model=shared.sd_model,
do_not_save_grid=True,
do_not_save_samples=True,
)
if preview_from_txt2img: if preview_from_txt2img:
p.prompt = preview_prompt p.prompt = preview_prompt
p.negative_prompt = preview_negative_prompt p.negative_prompt = preview_negative_prompt
p.steps = preview_steps p.steps = preview_steps
p.sampler_index = preview_sampler_index p.sampler_index = preview_sampler_index
p.cfg_scale = preview_cfg_scale p.cfg_scale = preview_cfg_scale
p.seed = preview_seed p.seed = preview_seed
p.width = preview_width p.width = preview_width
p.height = preview_height p.height = preview_height
else: else:
p.prompt = entries[0].cond_text p.prompt = entries[0].cond_text
p.steps = 20 p.steps = 20
preview_text = p.prompt preview_text = p.prompt
processed = processing.process_images(p) processed = processing.process_images(p)
image = processed.images[0] if len(processed.images)>0 else None image = processed.images[0] if len(processed.images)>0 else None
if unload: if unload:
shared.sd_model.cond_stage_model.to(devices.cpu) shared.sd_model.cond_stage_model.to(devices.cpu)
shared.sd_model.first_stage_model.to(devices.cpu) shared.sd_model.first_stage_model.to(devices.cpu)
if image is not None: if image is not None:
shared.state.current_image = image shared.state.current_image = image
last_saved_image, last_text_info = images.save_image(image, images_dir, "", p.seed, p.prompt, shared.opts.samples_format, processed.infotexts[0], p=p, forced_filename=forced_filename, save_to_dirs=False) last_saved_image, last_text_info = images.save_image(image, images_dir, "", p.seed, p.prompt, shared.opts.samples_format, processed.infotexts[0], p=p, forced_filename=forced_filename, save_to_dirs=False)
last_saved_image += f", prompt: {preview_text}" last_saved_image += f", prompt: {preview_text}"
shared.state.job_no = hypernetwork.step shared.state.job_no = hypernetwork.step
shared.state.textinfo = f""" shared.state.textinfo = f"""
<p> <p>
Loss: {previous_mean_loss:.7f}<br/> Loss: {previous_mean_loss:.7f}<br/>
Step: {hypernetwork.step}<br/> Step: {hypernetwork.step}<br/>
...@@ -510,7 +512,14 @@ Last saved hypernetwork: {html.escape(last_saved_file)}<br/> ...@@ -510,7 +512,14 @@ Last saved hypernetwork: {html.escape(last_saved_file)}<br/>
Last saved image: {html.escape(last_saved_image)}<br/> Last saved image: {html.escape(last_saved_image)}<br/>
</p> </p>
""" """
finally:
if weights:
for weight in weights:
weight.requires_grad = False
if unload:
shared.sd_model.cond_stage_model.to(devices.device)
shared.sd_model.first_stage_model.to(devices.device)
report_statistics(loss_dict) report_statistics(loss_dict)
checkpoint = sd_models.select_checkpoint() checkpoint = sd_models.select_checkpoint()
......
modules/textual_inversion/textual_inversion.py
View file @ 3ce2bfdf
...@@ -283,111 +283,113 @@ def train_embedding(embedding_name, learn_rate, batch_size, data_root, log_direc ...@@ -283,111 +283,113 @@ def train_embedding(embedding_name, learn_rate, batch_size, data_root, log_direc
embedding_yet_to_be_embedded = False embedding_yet_to_be_embedded = False
pbar = tqdm.tqdm(enumerate(ds), total=steps-ititial_step) pbar = tqdm.tqdm(enumerate(ds), total=steps-ititial_step)
for i, entries in pbar:
embedding.step = i + ititial_step
scheduler.apply(optimizer, embedding.step) try:
if scheduler.finished: for i, entries in pbar:
break embedding.step = i + ititial_step
if shared.state.interrupted: scheduler.apply(optimizer, embedding.step)
break if scheduler.finished:
break
with torch.autocast("cuda"):
c = cond_model([entry.cond_text for entry in entries]) if shared.state.interrupted:
x = torch.stack([entry.latent for entry in entries]).to(devices.device) break
loss = shared.sd_model(x, c)[0]
del x with torch.autocast("cuda"):
c = cond_model([entry.cond_text for entry in entries])
losses[embedding.step % losses.shape[0]] = loss.item() x = torch.stack([entry.latent for entry in entries]).to(devices.device)
loss = shared.sd_model(x, c)[0]
optimizer.zero_grad() del x
loss.backward()
optimizer.step() losses[embedding.step % losses.shape[0]] = loss.item()
steps_done = embedding.step + 1 optimizer.zero_grad()
loss.backward()
epoch_num = embedding.step // len(ds) optimizer.step()
epoch_step = embedding.step % len(ds)
steps_done = embedding.step + 1
pbar.set_description(f"[Epoch {epoch_num}: {epoch_step+1}/{len(ds)}]loss: {losses.mean():.7f}")
epoch_num = embedding.step // len(ds)
if embedding_dir is not None and steps_done % save_embedding_every == 0: epoch_step = embedding.step % len(ds)
# Before saving, change name to match current checkpoint.
embedding.name = f'{embedding_name}-{steps_done}' pbar.set_description(f"[Epoch {epoch_num}: {epoch_step+1}/{len(ds)}]loss: {losses.mean():.7f}")
last_saved_file = os.path.join(embedding_dir, f'{embedding.name}.pt')
embedding.save(last_saved_file) if embedding_dir is not None and steps_done % save_embedding_every == 0:
embedding_yet_to_be_embedded = True # Before saving, change name to match current checkpoint.
embedding.name = f'{embedding_name}-{steps_done}'
write_loss(log_directory, "textual_inversion_loss.csv", embedding.step, len(ds), { last_saved_file = os.path.join(embedding_dir, f'{embedding.name}.pt')
"loss": f"{losses.mean():.7f}", embedding.save(last_saved_file)
"learn_rate": scheduler.learn_rate embedding_yet_to_be_embedded = True
})
write_loss(log_directory, "textual_inversion_loss.csv", embedding.step, len(ds), {
if images_dir is not None and steps_done % create_image_every == 0: "loss": f"{losses.mean():.7f}",
forced_filename = f'{embedding_name}-{steps_done}' "learn_rate": scheduler.learn_rate
last_saved_image = os.path.join(images_dir, forced_filename) })
p = processing.StableDiffusionProcessingTxt2Img(
sd_model=shared.sd_model, if images_dir is not None and steps_done % create_image_every == 0:
do_not_save_grid=True, forced_filename = f'{embedding_name}-{steps_done}'
do_not_save_samples=True, last_saved_image = os.path.join(images_dir, forced_filename)
do_not_reload_embeddings=True, p = processing.StableDiffusionProcessingTxt2Img(
) sd_model=shared.sd_model,
do_not_save_grid=True,
if preview_from_txt2img: do_not_save_samples=True,
p.prompt = preview_prompt do_not_reload_embeddings=True,
p.negative_prompt = preview_negative_prompt )
p.steps = preview_steps
p.sampler_index = preview_sampler_index if preview_from_txt2img:
p.cfg_scale = preview_cfg_scale p.prompt = preview_prompt
p.seed = preview_seed p.negative_prompt = preview_negative_prompt
p.width = preview_width p.steps = preview_steps
p.height = preview_height p.sampler_index = preview_sampler_index
else: p.cfg_scale = preview_cfg_scale
p.prompt = entries[0].cond_text p.seed = preview_seed
p.steps = 20 p.width = preview_width
p.width = training_width p.height = preview_height
p.height = training_height else:
p.prompt = entries[0].cond_text
p.steps = 20
p.width = training_width
p.height = training_height
preview_text = p.prompt preview_text = p.prompt
processed = processing.process_images(p) processed = processing.process_images(p)
image = processed.images[0] image = processed.images[0]
shared.state.current_image = image shared.state.current_image = image
if save_image_with_stored_embedding and os.path.exists(last_saved_file) and embedding_yet_to_be_embedded: 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}-{steps_done}.png') last_saved_image_chunks = os.path.join(images_embeds_dir, f'{embedding_name}-{steps_done}.png')
info = PngImagePlugin.PngInfo() info = PngImagePlugin.PngInfo()
data = torch.load(last_saved_file) data = torch.load(last_saved_file)
info.add_text("sd-ti-embedding", embedding_to_b64(data)) info.add_text("sd-ti-embedding", embedding_to_b64(data))
title = "<{}>".format(data.get('name', '???')) title = "<{}>".format(data.get('name', '???'))
try: try:
vectorSize = list(data['string_to_param'].values())[0].shape[0] vectorSize = list(data['string_to_param'].values())[0].shape[0]
except Exception as e: except Exception as e:
vectorSize = '?' vectorSize = '?'
checkpoint = sd_models.select_checkpoint() checkpoint = sd_models.select_checkpoint()
footer_left = checkpoint.model_name footer_left = checkpoint.model_name
footer_mid = '[{}]'.format(checkpoint.hash) footer_mid = '[{}]'.format(checkpoint.hash)
footer_right = '{}v {}s'.format(vectorSize, steps_done) footer_right = '{}v {}s'.format(vectorSize, steps_done)
captioned_image = caption_image_overlay(image, title, footer_left, footer_mid, footer_right) captioned_image = caption_image_overlay(image, title, footer_left, footer_mid, footer_right)
captioned_image = insert_image_data_embed(captioned_image, data) captioned_image = insert_image_data_embed(captioned_image, data)
captioned_image.save(last_saved_image_chunks, "PNG", pnginfo=info) captioned_image.save(last_saved_image_chunks, "PNG", pnginfo=info)
embedding_yet_to_be_embedded = False embedding_yet_to_be_embedded = False
last_saved_image, last_text_info = images.save_image(image, images_dir, "", p.seed, p.prompt, shared.opts.samples_format, processed.infotexts[0], p=p, forced_filename=forced_filename, save_to_dirs=False) last_saved_image, last_text_info = images.save_image(image, images_dir, "", p.seed, p.prompt, shared.opts.samples_format, processed.infotexts[0], p=p, forced_filename=forced_filename, save_to_dirs=False)
last_saved_image += f", prompt: {preview_text}" last_saved_image += f", prompt: {preview_text}"
shared.state.job_no = embedding.step shared.state.job_no = embedding.step
shared.state.textinfo = f""" shared.state.textinfo = f"""
<p> <p>
Loss: {losses.mean():.7f}<br/> Loss: {losses.mean():.7f}<br/>
Step: {embedding.step}<br/> Step: {embedding.step}<br/>
...@@ -396,6 +398,9 @@ Last saved embedding: {html.escape(last_saved_file)}<br/> ...@@ -396,6 +398,9 @@ Last saved embedding: {html.escape(last_saved_file)}<br/>
Last saved image: {html.escape(last_saved_image)}<br/> Last saved image: {html.escape(last_saved_image)}<br/>
</p> </p>
""" """
finally:
if embedding and embedding.vec is not None:
embedding.vec.requires_grad = False
checkpoint = sd_models.select_checkpoint() checkpoint = sd_models.select_checkpoint()
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment