diff --git a/kohya_gui/class_tensorboard.py b/kohya_gui/class_tensorboard.py index 3c7480413..9fc1b76b5 100644 --- a/kohya_gui/class_tensorboard.py +++ b/kohya_gui/class_tensorboard.py @@ -113,4 +113,4 @@ def gradio_interface(self): self.stop_tensorboard, outputs=[button_start_tensorboard, button_stop_tensorboard], show_progress=False, - ) + ) \ No newline at end of file diff --git a/kohya_gui/dreambooth_gui.py b/kohya_gui/dreambooth_gui.py index 143a2c97b..4e9cf1499 100644 --- a/kohya_gui/dreambooth_gui.py +++ b/kohya_gui/dreambooth_gui.py @@ -510,12 +510,12 @@ def train_model( vae=vae, dataset_config=dataset_config, ): - return + return TRAIN_BUTTON_VISIBLE if not print_only and check_if_model_exist( output_name, output_dir, save_model_as, headless=headless ): - return + return TRAIN_BUTTON_VISIBLE if dataset_config: log.info( @@ -934,25 +934,13 @@ def dreambooth_tab( "Stop training", visible=False, variant="stop" ) - button_print = gr.Button("Print training command") + with gr.Column(), gr.Group(): + with gr.Row(): + button_print = gr.Button("Print training command") # Setup gradio tensorboard buttons with gr.Column(), gr.Group(): - ( - button_start_tensorboard, - button_stop_tensorboard, - ) = gradio_tensorboard() - - button_start_tensorboard.click( - start_tensorboard, - inputs=[dummy_headless, folders.logging_dir], - show_progress=False, - ) - - button_stop_tensorboard.click( - stop_tensorboard, - show_progress=False, - ) + TensorboardManager(headless=headless, logging_dir=folders.logging_dir) settings_list = [ source_model.pretrained_model_name_or_path, @@ -1142,4 +1130,4 @@ def dreambooth_tab( folders.reg_data_dir, folders.output_dir, folders.logging_dir, - ) + ) \ No newline at end of file diff --git a/kohya_gui/finetune_gui.py b/kohya_gui/finetune_gui.py index 2ed3370db..c15490cef 100644 --- a/kohya_gui/finetune_gui.py +++ b/kohya_gui/finetune_gui.py @@ -548,12 +548,12 @@ def train_model( resume=resume, dataset_config=dataset_config, ): - return + return TRAIN_BUTTON_VISIBLE if not print_only and check_if_model_exist( output_name, output_dir, save_model_as, headless ): - return + return TRAIN_BUTTON_VISIBLE if dataset_config: log.info( @@ -1086,24 +1086,13 @@ def list_presets(path): "Stop training", visible=False, variant="stop" ) - button_print = gr.Button("Print training command") + with gr.Column(), gr.Group(): + with gr.Row(): + button_print = gr.Button("Print training command") # Setup gradio tensorboard buttons with gr.Column(), gr.Group(): - ( - button_start_tensorboard, - button_stop_tensorboard, - ) = gradio_tensorboard() - - button_start_tensorboard.click( - start_tensorboard, - inputs=[dummy_headless, logging_dir], - ) - - button_stop_tensorboard.click( - stop_tensorboard, - show_progress=False, - ) + TensorboardManager(headless=headless, logging_dir=folders.logging_dir) settings_list = [ source_model.pretrained_model_name_or_path, @@ -1323,4 +1312,4 @@ def list_presets(path): if os.path.exists(top_level_path): with open(os.path.join(top_level_path), "r", encoding="utf8") as file: guides_top_level = file.read() + "\n" - gr.Markdown(guides_top_level) + gr.Markdown(guides_top_level) \ No newline at end of file diff --git a/kohya_gui/lora_gui.py b/kohya_gui/lora_gui.py index cd8928410..7aef5bab2 100644 --- a/kohya_gui/lora_gui.py +++ b/kohya_gui/lora_gui.py @@ -55,7 +55,7 @@ button_run = gr.Button("Start training", variant="primary") -button_stop_training = gr.Button("Stop training") +button_stop_training = gr.Button("Stop training", visible=False) document_symbol = "\U0001F4C4" # 📄 @@ -713,14 +713,14 @@ def train_model( lora_network_weights=lora_network_weights, dataset_config=dataset_config, ): - return + return TRAIN_BUTTON_VISIBLE if int(bucket_reso_steps) < 1: output_message( msg="Bucket resolution steps need to be greater than 0", headless=headless, ) - return + return TRAIN_BUTTON_VISIBLE # if noise_offset == "": # noise_offset = 0 @@ -730,7 +730,7 @@ def train_model( msg="Noise offset need to be a value between 0 and 1", headless=headless, ) - return + return TRAIN_BUTTON_VISIBLE if output_dir != "": if not os.path.exists(output_dir): @@ -746,7 +746,7 @@ def train_model( if not print_only and check_if_model_exist( output_name, output_dir, save_model_as, headless=headless ): - return + return TRAIN_BUTTON_VISIBLE # If string is empty set string to 0. # if text_encoder_lr == "": @@ -1015,7 +1015,7 @@ def train_model( # Sets flags for training specific components based on the provided learning rates. if float(learning_rate) == unet_lr_float == text_encoder_lr_float == 0: output_message(msg="Please input learning rate values.", headless=headless) - return + return TRAIN_BUTTON_VISIBLE # Flag to train text encoder only if its learning rate is non-zero and unet's is zero. network_train_text_encoder_only = text_encoder_lr_float != 0 and unet_lr_float == 0 # Flag to train unet only if its learning rate is non-zero and text encoder's is zero. @@ -2081,25 +2081,13 @@ def update_LoRA_settings( "Stop training", visible=False, variant="stop" ) - button_print = gr.Button("Print training command") + with gr.Column(), gr.Group(): + with gr.Row(): + button_print = gr.Button("Print training command") # Setup gradio tensorboard buttons with gr.Column(), gr.Group(): - ( - button_start_tensorboard, - button_stop_tensorboard, - ) = gradio_tensorboard() - - button_start_tensorboard.click( - start_tensorboard, - inputs=[dummy_headless, folders.logging_dir], - show_progress=False, - ) - - button_stop_tensorboard.click( - stop_tensorboard, - show_progress=False, - ) + TensorboardManager(headless=headless, logging_dir=folders.logging_dir) settings_list = [ source_model.pretrained_model_name_or_path, @@ -2352,4 +2340,4 @@ def update_LoRA_settings( folders.reg_data_dir, folders.output_dir, folders.logging_dir, - ) + ) \ No newline at end of file diff --git a/kohya_gui/tensorboard_gui.py b/kohya_gui/tensorboard_gui.py deleted file mode 100644 index 53b7116f5..000000000 --- a/kohya_gui/tensorboard_gui.py +++ /dev/null @@ -1,92 +0,0 @@ -import os -import gradio as gr -from easygui import msgbox -import subprocess -import time -import webbrowser -from .custom_logging import setup_logging - -# Set up logging -log = setup_logging() - -tensorboard_proc = None -TENSORBOARD = "tensorboard" if os.name == "posix" else "tensorboard.exe" - -# Set the default tensorboard port -DEFAULT_TENSORBOARD_PORT = 6006 - - -def start_tensorboard(headless, logging_dir, wait_time=5): - os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0" - global tensorboard_proc - - # Read the TENSORBOARD_PORT from the environment, or use the default - tensorboard_port = os.environ.get("TENSORBOARD_PORT", DEFAULT_TENSORBOARD_PORT) - - # Check if logging directory exists and is not empty; if not, warn the user and exit - if not os.path.exists(logging_dir) or not os.listdir(logging_dir): - log.error("Error: logging folder does not exist or does not contain logs.") - msgbox(msg="Error: logging folder does not exist or does not contain logs.") - return # Exit the function with an error code - - run_cmd = [ - TENSORBOARD, - "--logdir", - logging_dir, - "--host", - "0.0.0.0", - "--port", - str(tensorboard_port), - ] - - log.info(run_cmd) - if tensorboard_proc is not None: - log.info( - "Tensorboard is already running. Terminating existing process before starting new one..." - ) - stop_tensorboard() - - # Start background process - log.info("Starting TensorBoard on port {}".format(tensorboard_port)) - try: - # Copy the current environment - env = os.environ.copy() - - # Set your specific environment variable - env["TF_ENABLE_ONEDNN_OPTS"] = "0" - - tensorboard_proc = subprocess.Popen(run_cmd, env=env) - except Exception as e: - log.error("Failed to start Tensorboard:", e) - return - - if not headless: - # Wait for some time to allow TensorBoard to start up - time.sleep(wait_time) - - # Open the TensorBoard URL in the default browser - tensorboard_url = f"http://localhost:{tensorboard_port}" - log.info(f"Opening TensorBoard URL in browser: {tensorboard_url}") - webbrowser.open(tensorboard_url) - - -def stop_tensorboard(): - global tensorboard_proc - if tensorboard_proc is not None: - log.info("Stopping tensorboard process...") - try: - tensorboard_proc.terminate() - tensorboard_proc = None - log.info("...process stopped") - except Exception as e: - log.error("Failed to stop Tensorboard:", e) - else: - log.warning("Tensorboard is not running...") - - -def gradio_tensorboard(): - with gr.Row(): - button_start_tensorboard = gr.Button("Start tensorboard") - button_stop_tensorboard = gr.Button("Stop tensorboard") - - return (button_start_tensorboard, button_stop_tensorboard) diff --git a/kohya_gui/textual_inversion_gui.py b/kohya_gui/textual_inversion_gui.py index 7b9c9fe57..90f20367e 100644 --- a/kohya_gui/textual_inversion_gui.py +++ b/kohya_gui/textual_inversion_gui.py @@ -507,20 +507,20 @@ def train_model( vae=vae, dataset_config=dataset_config, ): - return + return TRAIN_BUTTON_VISIBLE if token_string == "": output_message(msg="Token string is missing", headless=headless) - return + return TRAIN_BUTTON_VISIBLE if init_word == "": output_message(msg="Init word is missing", headless=headless) - return + return TRAIN_BUTTON_VISIBLE if not print_only and check_if_model_exist( output_name, output_dir, save_model_as, headless ): - return + return TRAIN_BUTTON_VISIBLE if dataset_config: log.info( @@ -1046,25 +1046,13 @@ def list_embedding_files(path): "Stop training", visible=False, variant="stop" ) - button_print = gr.Button("Print training command") + with gr.Column(), gr.Group(): + with gr.Row(): + button_print = gr.Button("Print training command") # Setup gradio tensorboard buttons with gr.Column(), gr.Group(): - ( - button_start_tensorboard, - button_stop_tensorboard, - ) = gradio_tensorboard() - - button_start_tensorboard.click( - start_tensorboard, - inputs=[dummy_headless, folders.logging_dir], - show_progress=False, - ) - - button_stop_tensorboard.click( - stop_tensorboard, - show_progress=False, - ) + TensorboardManager(headless=headless, logging_dir=folders.logging_dir) settings_list = [ source_model.pretrained_model_name_or_path, @@ -1246,4 +1234,4 @@ def list_embedding_files(path): folders.reg_data_dir, folders.output_dir, folders.logging_dir, - ) + ) \ No newline at end of file diff --git a/requirements_windows.txt b/requirements_windows.txt index b5e37f79e..a9300090c 100644 --- a/requirements_windows.txt +++ b/requirements_windows.txt @@ -1,5 +1,5 @@ bitsandbytes==0.43.0 tensorboard -tensorflow +tensorflow>=2.16.1 onnxruntime-gpu==1.17.1 -r requirements.txt \ No newline at end of file