Merge pull request #2465 from bmaltais/dev

v24.1.2

.release

@@ -1 +1 @@
-v24.1.1
+v24.1.2

kohya_gui/basic_caption_gui.py

@@ -6,6 +6,7 @@
     find_replace,
     scriptdir,
     list_dirs,
+    setup_environment,
 )
 import os
 import sys
@@ -84,11 +85,7 @@ def caption_images(
         log.info(f"Executing command: {command_to_run}")
 
         # Set the environment variable for the Python path
-        env = os.environ.copy()
-        env["PYTHONPATH"] = (
-            rf"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}"
-        )
-        env["TF_ENABLE_ONEDNN_OPTS"] = "0"
+        env = setup_environment(scriptdir=scriptdir)
 
         # Run the command in the sd-scripts folder context
         subprocess.run(run_cmd, env=env, shell=False)

kohya_gui/blip_caption_gui.py

@@ -2,7 +2,7 @@
 import subprocess
 import os
 import sys
-from .common_gui import get_folder_path, add_pre_postfix, scriptdir, list_dirs
+from .common_gui import get_folder_path, add_pre_postfix, scriptdir, list_dirs, setup_environment
 from .custom_logging import setup_logging
 
 # Set up logging
@@ -87,11 +87,7 @@ def caption_images(
     )
 
     # Set up the environment
-    env = os.environ.copy()
-    env["PYTHONPATH"] = (
-        rf"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}"
-    )
-    env["TF_ENABLE_ONEDNN_OPTS"] = "0"
+    env = setup_environment(scriptdir=scriptdir)
 
     # Reconstruct the safe command string for display
     command_to_run = " ".join(run_cmd)

kohya_gui/class_tensorboard.py

@@ -15,6 +15,7 @@
 from easygui import msgbox
 from threading import Thread, Event
 from .custom_logging import setup_logging
+from .common_gui import setup_environment
 
 
 class TensorboardManager:
@@ -72,8 +73,7 @@ def start_tensorboard(self, logging_dir=None):
 
         self.log.info("Starting TensorBoard on port {}".format(self.tensorboard_port))
         try:
-            env = os.environ.copy()
-            env["TF_ENABLE_ONEDNN_OPTS"] = "0"
+            env = setup_environment(scriptdir=scriptdir)
             self.tensorboard_proc = subprocess.Popen(run_cmd, env=env)
         except Exception as e:
             self.log.error("Failed to start Tensorboard:", e)

kohya_gui/common_gui.py

@@ -1420,10 +1420,8 @@ def validate_model_path(pretrained_model_name_or_path: str) -> bool:
         log.info(f"{msg} SUCCESS")
     else:
         # If not one of the default models, check if it's a valid local path
-        if not os.path.exists(pretrained_model_name_or_path):
-            log.error(f"{msg} FAILED: is missing or does not exist")
+        if not validate_file_path(pretrained_model_name_or_path):
             return False
-        log.info(f"{msg} SUCCESS")
 
     return True
 
@@ -1486,3 +1484,15 @@ def validate_args_setting(input_string):
             "A valid settings string must consist of one or more key/value pairs formatted as key=value, with no spaces around the equals sign or within the value. Multiple pairs should be separated by a space."
         )
         return False
+
+def setup_environment(scriptdir: str):
+    env = os.environ.copy()
+    env["PYTHONPATH"] = (
+        fr"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}"
+    )
+    env["TF_ENABLE_ONEDNN_OPTS"] = "0"
+
+    if os.name == "nt":
+        env["XFORMERS_FORCE_DISABLE_TRITON"] = "1"
+
+    return env

kohya_gui/convert_lcm_gui.py

@@ -7,7 +7,7 @@
     get_file_path,
     scriptdir,
     list_files,
-    create_refresh_button,
+    create_refresh_button, setup_environment
 )
 from .custom_logging import setup_logging
 
@@ -62,11 +62,7 @@ def convert_lcm(
         run_cmd.append("--ssd-1b")
 
     # Set up the environment
-    env = os.environ.copy()
-    env["PYTHONPATH"] = (
-        fr"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}"
-    )
-    env["TF_ENABLE_ONEDNN_OPTS"] = "0"
+    env = setup_environment(scriptdir=scriptdir)
 
     # Reconstruct the safe command string for display
     command_to_run = " ".join(run_cmd)

kohya_gui/convert_model_gui.py

@@ -2,7 +2,7 @@
 import subprocess
 import os
 import sys
-from .common_gui import get_folder_path, get_file_path, scriptdir, list_files, list_dirs
+from .common_gui import get_folder_path, get_file_path, scriptdir, list_files, list_dirs, setup_environment
 
 from .custom_logging import setup_logging
 
@@ -99,12 +99,7 @@ def convert_model(
     # Log the command
     log.info(" ".join(run_cmd))
 
-    env = os.environ.copy()
-    env["PYTHONPATH"] = (
-        rf"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}"
-    )
-    # Adding an example of an environment variable that might be relevant
-    env["TF_ENABLE_ONEDNN_OPTS"] = "0"
+    env = setup_environment(scriptdir=scriptdir)
 
     # Run the command
     subprocess.run(run_cmd, env=env, shell=False)

kohya_gui/dreambooth_gui.py

@@ -19,6 +19,7 @@
     update_my_data,
     validate_file_path, validate_folder_path, validate_model_path,
     validate_args_setting,
+    setup_environment,
 )
 from .class_accelerate_launch import AccelerateLaunch
 from .class_configuration_file import ConfigurationFile
@@ -541,7 +542,7 @@ def train_model(
     if not validate_folder_path(train_data_dir):
         return TRAIN_BUTTON_VISIBLE
 
-    if not validate_file_path(vae):
+    if not validate_model_path(vae):
         return TRAIN_BUTTON_VISIBLE
     #
     # End of path validation
@@ -898,11 +899,7 @@ def train_model(
 
         # log.info(run_cmd)
 
-        env = os.environ.copy()
-        env["PYTHONPATH"] = (
-            fr"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}"
-        )
-        env["TF_ENABLE_ONEDNN_OPTS"] = "0"
+        env = setup_environment(scriptdir=scriptdir)
 
         # Run the command
 

kohya_gui/extract_lora_from_dylora_gui.py

@@ -6,7 +6,7 @@
     get_file_path,
     scriptdir,
     list_files,
-    create_refresh_button,
+    create_refresh_button, setup_environment
 )
 
 from .custom_logging import setup_logging
@@ -59,12 +59,7 @@ def extract_dylora(
         str(unit),
     ]
 
-    env = os.environ.copy()
-    env["PYTHONPATH"] = (
-        fr"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}"
-    )
-    # Example environment variable adjustment for the Python environment
-    env["TF_ENABLE_ONEDNN_OPTS"] = "0"
+    env = setup_environment(scriptdir=scriptdir)
 
     # Reconstruct the safe command string for display
     command_to_run = " ".join(run_cmd)

kohya_gui/extract_lora_gui.py

@@ -8,7 +8,7 @@
     is_file_writable,
     scriptdir,
     list_files,
-    create_refresh_button,
+    create_refresh_button, setup_environment
 )
 
 from .custom_logging import setup_logging
@@ -109,12 +109,7 @@ def extract_lora(
         run_cmd.append("--load_tuned_model_to")
         run_cmd.append(load_tuned_model_to)
 
-    env = os.environ.copy()
-    env["PYTHONPATH"] = (
-        rf"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}"
-    )
-    # Adding an example of another potentially relevant environment variable
-    env["TF_ENABLE_ONEDNN_OPTS"] = "0"
+    env = setup_environment(scriptdir=scriptdir)
 
     # Reconstruct the safe command string for display
     command_to_run = " ".join(run_cmd)

kohya_gui/extract_lycoris_locon_gui.py

@@ -7,7 +7,7 @@
     get_file_path,
     scriptdir,
     list_files,
-    create_refresh_button,
+    create_refresh_button, setup_environment
 )
 
 from .custom_logging import setup_logging
@@ -123,12 +123,7 @@ def extract_lycoris_locon(
     run_cmd.append(fr"{db_model}")
     run_cmd.append(fr"{output_name}")
 
-    env = os.environ.copy()
-    env["PYTHONPATH"] = (
-        fr"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}"
-    )
-    # Adding an example of an environment variable that might be relevant
-    env["TF_ENABLE_ONEDNN_OPTS"] = "0"
+    env = setup_environment(scriptdir=scriptdir)
 
     # Reconstruct the safe command string for display
     command_to_run = " ".join(run_cmd)

kohya_gui/finetune_gui.py

@@ -19,7 +19,7 @@
     scriptdir,
     update_my_data,
     validate_file_path, validate_folder_path, validate_model_path,
-    validate_args_setting,
+    validate_args_setting, setup_environment,
 )
 from .class_accelerate_launch import AccelerateLaunch
 from .class_configuration_file import ConfigurationFile
@@ -639,15 +639,7 @@ def train_model(
             log.info(" ".join(run_cmd))
 
             # Prepare environment variables
-            env = os.environ.copy()
-            env["PYTHONPATH"] = (
-                rf"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}"
-            )
-            env["TF_ENABLE_ONEDNN_OPTS"] = "0"
-
-            # Execute the command if not in print-only mode
-            if not print_only:
-                subprocess.run(run_cmd, env=env)
+            env = setup_environment(scriptdir=scriptdir)
 
         # create images buckets
         if generate_image_buckets:
@@ -685,11 +677,7 @@ def train_model(
             log.info(" ".join(run_cmd))
 
             # Copy and modify environment variables
-            env = os.environ.copy()
-            env["PYTHONPATH"] = (
-                rf"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}"
-            )
-            env["TF_ENABLE_ONEDNN_OPTS"] = "0"
+            env = setup_environment(scriptdir=scriptdir)
 
             # Execute the command if not just for printing
             if not print_only:
@@ -965,11 +953,7 @@ def train_model(
 
         # log.info(run_cmd)
 
-        env = os.environ.copy()
-        env["PYTHONPATH"] = (
-            rf"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}"
-        )
-        env["TF_ENABLE_ONEDNN_OPTS"] = "0"
+        env = setup_environment(scriptdir=scriptdir)
 
         # Run the command
         executor.execute_command(run_cmd=run_cmd, env=env)

kohya_gui/git_caption_gui.py

@@ -2,7 +2,7 @@
 import subprocess
 import os
 import sys
-from .common_gui import get_folder_path, add_pre_postfix, scriptdir, list_dirs
+from .common_gui import get_folder_path, add_pre_postfix, scriptdir, list_dirs, setup_environment
 
 from .custom_logging import setup_logging
 
@@ -58,12 +58,7 @@ def caption_images(
     # Add the directory containing the training data
     run_cmd.append(fr"{train_data_dir}")
 
-    env = os.environ.copy()
-    env["PYTHONPATH"] = (
-        fr"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}"
-    )
-    # Adding an example of an environment variable that might be relevant
-    env["TF_ENABLE_ONEDNN_OPTS"] = "0"
+    env = setup_environment(scriptdir=scriptdir)
 
     # Reconstruct the safe command string for display
     command_to_run = " ".join(run_cmd)

kohya_gui/group_images_gui.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import subprocess
-from .common_gui import get_folder_path, scriptdir, list_dirs
+from .common_gui import get_folder_path, scriptdir, list_dirs, setup_environment
 import os
 import sys
 
@@ -51,12 +51,7 @@ def group_images(
             run_cmd.append("--caption_ext")
             run_cmd.append(caption_ext)
 
-    env = os.environ.copy()
-    env["PYTHONPATH"] = (
-        rf"{scriptdir}{os.pathsep}{scriptdir}/tools{os.pathsep}{env.get('PYTHONPATH', '')}"
-    )
-    # Adding a common environmental setting as an example if it's missing in the original context
-    env["TF_ENABLE_ONEDNN_OPTS"] = "0"
+    env = setup_environment(scriptdir=scriptdir)
 
     # Reconstruct the safe command string for display
     command_to_run = " ".join(run_cmd)

kohya_gui/lora_gui.py

@@ -20,7 +20,7 @@
     scriptdir,
     update_my_data,
     validate_file_path, validate_folder_path, validate_model_path, validate_toml_file,
-    validate_args_setting,
+    validate_args_setting, setup_environment,
 )
 from .class_accelerate_launch import AccelerateLaunch
 from .class_configuration_file import ConfigurationFile
@@ -732,7 +732,7 @@ def train_model(
     if not validate_folder_path(train_data_dir):
         return TRAIN_BUTTON_VISIBLE
 
-    if not validate_file_path(vae):
+    if not validate_model_path(vae):
         return TRAIN_BUTTON_VISIBLE
 
     #
@@ -1282,11 +1282,7 @@ def train_model(
         )
 
         # log.info(run_cmd)
-        env = os.environ.copy()
-        env["PYTHONPATH"] = (
-            rf"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}"
-        )
-        env["TF_ENABLE_ONEDNN_OPTS"] = "0"
+        env = setup_environment(scriptdir=scriptdir)
 
         # Run the command
 

kohya_gui/merge_lora_gui.py

@@ -13,7 +13,7 @@
     get_file_path,
     scriptdir,
     list_files,
-    create_refresh_button,
+    create_refresh_button, setup_environment
 )
 from .custom_logging import setup_logging
 
@@ -451,12 +451,7 @@ def merge_lora(
                 map(str, valid_ratios)
             )  # Convert ratios to strings and include them as separate arguments
 
-        env = os.environ.copy()
-        env["PYTHONPATH"] = (
-            rf"{scriptdir}{os.pathsep}{scriptdir}/sd-scripts{os.pathsep}{env.get('PYTHONPATH', '')}"
-        )
-        # Example of adding an environment variable for TensorFlow, if necessary
-        env["TF_ENABLE_ONEDNN_OPTS"] = "0"
+        env = setup_environment(scriptdir=scriptdir)
 
         # Reconstruct the safe command string for display
         command_to_run = " ".join(run_cmd)