From 0d4e6c4ece67e8375062cf39f3a8a1e83317382f Mon Sep 17 00:00:00 2001 From: bmaltais Date: Sat, 16 Mar 2024 11:00:13 -0400 Subject: [PATCH] Update command executor code --- kohya_gui/class_command_executor.py | 35 +++++++++++++++++++++-------- kohya_gui/class_lora_tab.py | 10 ++++----- kohya_gui/lora_gui.py | 8 +------ 3 files changed, 32 insertions(+), 21 deletions(-) diff --git a/kohya_gui/class_command_executor.py b/kohya_gui/class_command_executor.py index c7b55e5..d64ad92 100644 --- a/kohya_gui/class_command_executor.py +++ b/kohya_gui/class_command_executor.py @@ -5,30 +5,47 @@ from .custom_logging import setup_logging # Set up logging log = setup_logging() - class CommandExecutor: + """ + A class to execute and manage commands. + """ + def __init__(self): + """ + Initialize the CommandExecutor. + """ self.process = None - def execute_command(self, run_cmd, **kwargs): + def execute_command(self, run_cmd: str, **kwargs): + """ + Execute a command if no other command is currently running. + + Parameters: + - run_cmd (str): The command to execute. + - **kwargs: Additional keyword arguments to pass to subprocess.Popen. + """ if self.process and self.process.poll() is None: - log.info( - 'The command is already running. Please wait for it to finish.' - ) + log.info("The command is already running. Please wait for it to finish.") else: self.process = subprocess.Popen(run_cmd, shell=True, **kwargs) def kill_command(self): + """ + Kill the currently running command and its child processes. + """ if self.process and self.process.poll() is None: try: + # Get the parent process and kill all its children parent = psutil.Process(self.process.pid) for child in parent.children(recursive=True): child.kill() parent.kill() - log.info('The running process has been terminated.') + log.info("The running process has been terminated.") except psutil.NoSuchProcess: - log.info('The process does not exist.') + # Explicitly handle the case where the process does not exist + log.info("The process does not exist. It might have terminated before the kill command was issued.") except Exception as e: - log.info(f'Error when terminating process: {e}') + # General exception handling for any other errors + log.info(f"Error when terminating process: {e}") else: - log.info('There is no running process to kill.') + log.info("There is no running process to kill.") diff --git a/kohya_gui/class_lora_tab.py b/kohya_gui/class_lora_tab.py index bcc1d16..d179307 100644 --- a/kohya_gui/class_lora_tab.py +++ b/kohya_gui/class_lora_tab.py @@ -10,14 +10,14 @@ from .extract_lora_from_dylora_gui import gradio_extract_dylora_tab from .merge_lycoris_gui import gradio_merge_lycoris_tab # Deprecated code -from .dataset_balancing_gui import gradio_dataset_balancing_tab -from .dreambooth_folder_creation_gui import ( - gradio_dreambooth_folder_creation_tab, -) +# from .dataset_balancing_gui import gradio_dataset_balancing_tab +# from .dreambooth_folder_creation_gui import ( +# gradio_dreambooth_folder_creation_tab, +# ) class LoRATools: - def __init__(self, train_data_dir=None, reg_data_dir=None, output_dir=None, logging_dir=None, headless: bool = False): + def __init__(self, headless: bool = False): self.headless = headless gr.Markdown( diff --git a/kohya_gui/lora_gui.py b/kohya_gui/lora_gui.py index a953239..1ab1ee7 100644 --- a/kohya_gui/lora_gui.py +++ b/kohya_gui/lora_gui.py @@ -1988,13 +1988,7 @@ def lora_tab( ) with gr.Tab("Tools"): - lora_tools = LoRATools( - train_data_dir=source_model.train_data_dir, - reg_data_dir=folders.reg_data_dir, - output_dir=folders.output_dir, - logging_dir=folders.logging_dir, - headless=headless - ) + lora_tools = LoRATools(headless=headless) with gr.Tab("Guides"): gr.Markdown("This section provide Various LoRA guides and information...")