OPENGIT/stable-diffusion-webui
1
0
Fork 0
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui.git synced 2025-04-13 00:09:02 +08:00
Code Issues Packages Projects Releases Wiki Activity

Make Tensorboard a late import (it was implicitly installed by basicsr)

Browse Source
This commit is contained in:
Aarni Koskela 2023-12-30 19:44:05 +02:00
parent c9174253fb
commit 1465dab715
1 changed files with 7 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
modules/textual_inversion/textual_inversion.py
View File

@ -11,7 +11,6 @@ import safetensors.torch
import numpy as np import numpy as np
from PIL import Image, PngImagePlugin from PIL import Image, PngImagePlugin
from torch.utils.tensorboard import SummaryWriter
from modules import shared, devices, sd_hijack, sd_models, images, sd_samplers, sd_hijack_checkpoint, errors, hashes from modules import shared, devices, sd_hijack, sd_models, images, sd_samplers, sd_hijack_checkpoint, errors, hashes
import modules.textual_inversion.dataset import modules.textual_inversion.dataset
@ -344,6 +343,7 @@ def write_loss(log_directory, filename, step, epoch_len, values):
}) })
def tensorboard_setup(log_directory): def tensorboard_setup(log_directory):
from torch.utils.tensorboard import SummaryWriter
os.makedirs(os.path.join(log_directory, "tensorboard"), exist_ok=True) os.makedirs(os.path.join(log_directory, "tensorboard"), exist_ok=True)
return SummaryWriter( return SummaryWriter(
log_dir=os.path.join(log_directory, "tensorboard"), log_dir=os.path.join(log_directory, "tensorboard"),
@ -448,8 +448,12 @@ def train_embedding(id_task, embedding_name, learn_rate, batch_size, gradient_st
shared.state.textinfo = f"Preparing dataset from {html.escape(data_root)}..." shared.state.textinfo = f"Preparing dataset from {html.escape(data_root)}..."
old_parallel_processing_allowed = shared.parallel_processing_allowed old_parallel_processing_allowed = shared.parallel_processing_allowed
tensorboard_writer = None
if shared.opts.training_enable_tensorboard: if shared.opts.training_enable_tensorboard:
try:
tensorboard_writer = tensorboard_setup(log_directory) tensorboard_writer = tensorboard_setup(log_directory)
except ImportError:
errors.report("Error initializing tensorboard", exc_info=True)
pin_memory = shared.opts.pin_memory pin_memory = shared.opts.pin_memory
@ -622,7 +626,7 @@ def train_embedding(id_task, embedding_name, learn_rate, batch_size, gradient_st
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}"
if shared.opts.training_enable_tensorboard and shared.opts.training_tensorboard_save_images: if tensorboard_writer and shared.opts.training_tensorboard_save_images:
tensorboard_add_image(tensorboard_writer, f"Validation at epoch {epoch_num}", image, embedding.step) tensorboard_add_image(tensorboard_writer, f"Validation at epoch {epoch_num}", image, embedding.step)
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: