diff --git a/README.md b/README.md index 252653a..7d8182d 100644 --- a/README.md +++ b/README.md @@ -27,7 +27,8 @@ Support for Linux and macOS is also available. While Linux support is actively m - [Colab](#-colab) - [Runpod, Novita, Docker](#runpod-novita-docker) - [Custom Path Defaults](#custom-path-defaults) - - [LoRA](#lora) +- [Server Options](#server-options) +- [LoRA](#lora) - [Sample image generation during training](#sample-image-generation-during-training) - [Troubleshooting](#troubleshooting) - [Page File Limit](#page-file-limit) @@ -101,9 +102,9 @@ I would like to express my gratitude to camenduru for their valuable contributio These options are for users running training on hosted GPU infrastructure or containers. -- **[Runpod setup](docs/runpod_setup.md)** – Ready-made GPU background training via templates. -- **[Novita setup](docs/novita_setup.md)** – Similar to Runpod, but integrated into the Novita UI. -- **[Docker setup](docs/docker.md)** – For developers/sysadmins using containerized environments. +- **[Runpod setup](docs/installation_runpod.md)** – Ready-made GPU background training via templates. +- **[Novita setup](docs/installation_novita.md)** – Similar to Runpod, but integrated into the Novita UI. +- **[Docker setup](docs/installation_docker.md)** – For developers/sysadmins using containerized environments. ## Custom Path Defaults with `config.toml` @@ -180,6 +181,21 @@ If you prefer to name your configuration file differently or store it in another By effectively using `config.toml`, you can significantly speed up your training setup process. Always refer to the `config example.toml` for the most up-to-date list of configurable paths. +## Server Options + +The `config.toml` file can also be used to configure the Gradio server. + +### `allowed_paths` + +The `allowed_paths` option allows you to specify a list of directories that the application can access. This is useful if you want to store your models, datasets, or other files on an external drive or in a location outside of the application's root directory. + +**Example:** + +```toml +[server] +allowed_paths = ["/mnt/external_drive/models", "/home/user/datasets"] +``` + ## LoRA To train a LoRA, you can currently use the `train_network.py` code. You can create a LoRA network by using the all-in-one GUI. diff --git a/config example.toml b/config example.toml index 3ab9609..cbfc760 100644 --- a/config example.toml +++ b/config example.toml @@ -4,6 +4,9 @@ [settings] use_shell = false # Use shell during process run of sd-scripts oython code. Most secure is false but some systems may require it to be true to properly run sd-scripts. +[server] +allowed_paths = [] # A list of paths that the application is allowed to access. + # Default folders location [model] models_dir = "./models" # Pretrained model name or path diff --git a/kohya_gui.py b/kohya_gui.py index e8d64a8..0954918 100644 --- a/kohya_gui.py +++ b/kohya_gui.py @@ -102,6 +102,7 @@ def UI(**kwargs): "share": False if kwargs.get("do_not_share", False) else kwargs.get("share", False), "root_path": kwargs.get("root_path", None), "debug": kwargs.get("debug", False), + "allowed_paths": config.allowed_paths, } # This line filters out any key-value pairs from `launch_params` where the value is `None`, ensuring only valid parameters are passed to the `launch` function. diff --git a/kohya_gui/class_gui_config.py b/kohya_gui/class_gui_config.py index 33064a9..92c20da 100644 --- a/kohya_gui/class_gui_config.py +++ b/kohya_gui/class_gui_config.py @@ -16,6 +16,7 @@ class KohyaSSGUIConfig: Initialize the KohyaSSGUIConfig class. """ self.config = self.load_config(config_file_path=config_file_path) + self.allowed_paths = self.get_allowed_paths() def load_config(self, config_file_path: str = "./config.toml") -> dict: """ @@ -81,6 +82,12 @@ class KohyaSSGUIConfig: log.debug(f"Returned {data}") return data + def get_allowed_paths(self) -> list: + """ + Retrieves the list of allowed paths from the [server] section of the config file. + """ + return self.get("server.allowed_paths", []) + def is_config_loaded(self) -> bool: """ Checks if the configuration was loaded from a file. diff --git a/test/test_allowed_paths.py b/test/test_allowed_paths.py new file mode 100644 index 0000000..a99b8a7 --- /dev/null +++ b/test/test_allowed_paths.py @@ -0,0 +1,59 @@ +import os +import tempfile +import unittest +import subprocess +from pathlib import Path + +class TestAllowedPaths(unittest.TestCase): + def setUp(self): + print("Setting up test...") + self.temp_dir = tempfile.TemporaryDirectory() + self.temp_dir_path = Path(self.temp_dir.name) + self.dummy_file = self.temp_dir_path / "dummy.txt" + with open(self.dummy_file, "w") as f: + f.write("dummy content") + + self.config_file = self.temp_dir_path / "config.toml" + with open(self.config_file, "w") as f: + f.write(f''' +[server] +allowed_paths = ["{self.temp_dir.name}"] +''') + print("Setup complete.") + + def tearDown(self): + print("Tearing down test...") + self.temp_dir.cleanup() + print("Teardown complete.") + + def test_allowed_paths(self): + print("Running test_allowed_paths...") + # Run the gui with the new config and check if it can access the dummy file + process = subprocess.Popen( + [ + "python", + "kohya_gui.py", + "--config", + str(self.config_file), + "--headless", + ], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + + print("Process started.") + # Give the server some time to start + try: + stdout, stderr = process.communicate(timeout=10) + except subprocess.TimeoutExpired: + process.kill() + stdout, stderr = process.communicate() + + print(f"Stdout: {stdout.decode()}") + print(f"Stderr: {stderr.decode()}") + # Check if there are any errors in the stderr + self.assertNotIn("InvalidPathError", stderr.decode()) + print("Test complete.") + +if __name__ == "__main__": + unittest.main()