From ef6feedeb22c6c3d045f491200bd7237735b9f78 Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Sat, 18 Nov 2023 23:38:39 -0300
Subject: [PATCH] Add --nowebui flag for pure API mode (#4651)

---
 README.md                   |  1 +
 extensions/openai/script.py |  5 ++++-
 modules/shared.py           |  3 ++-
 server.py                   | 24 +++++++++++++++---------
 4 files changed, 22 insertions(+), 11 deletions(-)

diff --git a/README.md b/README.md
index 8c2679cf..56e810b2 100644
--- a/README.md
+++ b/README.md
@@ -414,6 +414,7 @@ Optionally, you can use the following command-line flags:
 | `--api-port API_PORT`                 | The listening port for the API. |
 | `--api-key API_KEY`                   | API authentication key. |
 | `--admin-key ADMIN_KEY`               | API authentication key for admin tasks like loading and unloading models. If not set, will be the same as --api-key. |
+| `--nowebui`                           | Do not launch the Gradio UI. Useful for launching the API in standalone mode. |
 
 #### Multimodal
 
diff --git a/extensions/openai/script.py b/extensions/openai/script.py
index da56287c..a516b0f7 100644
--- a/extensions/openai/script.py
+++ b/extensions/openai/script.py
@@ -341,4 +341,7 @@ def run_server():
 
 
 def setup():
-    Thread(target=run_server, daemon=True).start()
+    if shared.args.nowebui:
+        run_server()
+    else:
+        Thread(target=run_server, daemon=True).start()
diff --git a/modules/shared.py b/modules/shared.py
index b139a2cf..344daf1d 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -171,6 +171,7 @@ parser.add_argument('--public-api-id', type=str, help='Tunnel ID for named Cloud
 parser.add_argument('--api-port', type=int, default=5000, help='The listening port for the API.')
 parser.add_argument('--api-key', type=str, default='', help='API authentication key.')
 parser.add_argument('--admin-key', type=str, default='', help='API authentication key for admin tasks like loading and unloading models. If not set, will be the same as --api-key.')
+parser.add_argument('--nowebui', action='store_true', help='Do not launch the Gradio UI. Useful for launching the API in standalone mode.')
 
 # Multimodal
 parser.add_argument('--multimodal-pipeline', type=str, default=None, help='The multimodal pipeline to use. Examples: llava-7b, llava-13b.')
@@ -201,7 +202,7 @@ for k in ['notebook', 'chat', 'no_stream', 'mul_mat_q', 'use_fast']:
 # Security warnings
 if args.trust_remote_code:
     logger.warning('trust_remote_code is enabled. This is dangerous.')
-if 'COLAB_GPU' not in os.environ:
+if 'COLAB_GPU' not in os.environ and not args.nowebui:
     if args.share:
         logger.warning("The gradio \"share link\" feature uses a proprietary executable to create a reverse tunnel. Use it with care.")
     if any((args.listen, args.share)) and not any((args.gradio_auth, args.gradio_auth_path)):
diff --git a/server.py b/server.py
index e9605e3b..cdd82e1d 100644
--- a/server.py
+++ b/server.py
@@ -226,13 +226,19 @@ if __name__ == "__main__":
 
     shared.generation_lock = Lock()
 
-    # Launch the web UI
-    create_interface()
-    while True:
-        time.sleep(0.5)
-        if shared.need_restart:
-            shared.need_restart = False
+    if shared.args.nowebui:
+        # Start the API in standalone mode
+        shared.args.extensions = [x for x in shared.args.extensions if x != 'gallery']
+        if shared.args.extensions is not None and len(shared.args.extensions) > 0:
+            extensions_module.load_extensions()
+    else:
+        # Launch the web UI
+        create_interface()
+        while True:
             time.sleep(0.5)
-            shared.gradio['interface'].close()
-            time.sleep(0.5)
-            create_interface()
+            if shared.need_restart:
+                shared.need_restart = False
+                time.sleep(0.5)
+                shared.gradio['interface'].close()
+                time.sleep(0.5)
+                create_interface()