From a3275fc0dca7c4a862b362175465de5a3f19865e Mon Sep 17 00:00:00 2001
From: Qing <cwq1913@gmail.com>
Date: Tue, 17 Jan 2023 20:48:38 +0800
Subject: [PATCH] update user scripts

---
 scripts/user_scripts/tasks.py       | 27 +++++++++++++++++++--------
 scripts/user_scripts/win_config.bat |  2 ++
 scripts/user_scripts/win_start.bat  |  4 ++++
 3 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/scripts/user_scripts/tasks.py b/scripts/user_scripts/tasks.py
index 2ed405d..f77649f 100644
--- a/scripts/user_scripts/tasks.py
+++ b/scripts/user_scripts/tasks.py
@@ -105,12 +105,23 @@ def start(c):
     port = find_free_port()
     log.info(f"Using random port: {port}")
 
-    if desktop:
-        c.run(
-            f"lama-cleaner --model {model} --device {device} --port {port} --gui --gui-size 1400 900"
-        )
-    else:
-        c.run(
-            f"lama-cleaner --model {model} --device {device} --port {port}"
-        )
+    commandline_args = [
+        "--model", model,
+        "--device", device,
+        "--port", port,
+    ]
+
+    if desktop:
+        commandline_args.extend(["--gui", "--gui-size", "1400", "900"])
+
+    model_dir = os.environ.get('MODEL_DIR', "")
+    if model_dir:
+        commandline_args.extend(["--model-dir", model_dir])
+    
+    commandline_args = ' '.join(commandline_args)
+    env_commandline_args = os.environ.get('COMMANDLINE_ARGS', "")
+
+    c.run(
+        f"lama-cleaner {env_commandline_args} {commandline_args}"
+    )
 
diff --git a/scripts/user_scripts/win_config.bat b/scripts/user_scripts/win_config.bat
index 19ab7d6..9628219 100644
--- a/scripts/user_scripts/win_config.bat
+++ b/scripts/user_scripts/win_config.bat
@@ -12,3 +12,5 @@ set PATH=C:\Windows\System32;%PATH%
 
 
 @call invoke config
+
+PAUSE
\ No newline at end of file
diff --git a/scripts/user_scripts/win_start.bat b/scripts/user_scripts/win_start.bat
index c9b1648..b4a094a 100644
--- a/scripts/user_scripts/win_start.bat
+++ b/scripts/user_scripts/win_start.bat
@@ -4,4 +4,8 @@ set PATH=C:\Windows\System32;%PATH%
 
 @call installer\Scripts\activate.bat
 
+set MODEL_DIR=
+set COMMANDLINE_ARGS=
 @call invoke start
+
+PAUSE
\ No newline at end of file