1 files changed, 23 insertions, 1 deletions
diff --git a/modules/paths.py b/modules/paths.py
index 4dd03a35..d991cc71 100644
--- a/modules/paths.py
+++ b/modules/paths.py
@@ -4,7 +4,15 @@ import sys
 import modules.safe
 
 script_path = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
-models_path = os.path.join(script_path, "models")
+
+# Parse the --data-dir flag first so we can use it as a base for our other argument default values
+parser = argparse.ArgumentParser(add_help=False)
+parser.add_argument("--data-dir", type=str, default=os.path.dirname(os.path.dirname(os.path.realpath(__file__))), help="base path where all user data is stored",)
+cmd_opts_pre = parser.parse_known_args()[0]
+data_path = cmd_opts_pre.data_dir
+models_path = os.path.join(data_path, "models")
+
+# data_path = cmd_opts_pre.data
 sys.path.insert(0, script_path)
 
 # search for directory of stable diffusion in following places
@@ -38,3 +46,17 @@ for d, must_exist, what, options in path_dirs:
         else:
             sys.path.append(d)
         paths[what] = d
+
+
+class Prioritize:
+    def __init__(self, name):
+        self.name = name
+        self.path = None
+
+    def __enter__(self):
+        self.path = sys.path.copy()
+        sys.path = [paths[self.name]] + sys.path
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        sys.path = self.path
+        self.path = None