Add CORS-allow policy launch argument using regex

2 files changed, 9 insertions, 4 deletions

diff --git a/modules/shared.py b/modules/shared.py
index e8bacd3c..55de286d 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -81,12 +81,13 @@ parser.add_argument("--disable-console-progressbars", action='store_true', help=
 parser.add_argument("--enable-console-prompts", action='store_true', help="print prompts to console when generating with txt2img and img2img", default=False)

 parser.add_argument('--vae-path', type=str, help='Path to Variational Autoencoders model', default=None)

 parser.add_argument("--disable-safe-unpickle", action='store_true', help="disable checking pytorch models for malicious code", default=False)

-parser.add_argument("--api", action='store_true', help="use api=True to launch the api with the webui")

-parser.add_argument("--nowebui", action='store_true', help="use api=True to launch the api instead of the webui")

+parser.add_argument("--api", action='store_true', help="use api=True to launch the API together with the webui (use --nowebui instead for only the API)")

+parser.add_argument("--nowebui", action='store_true', help="use api=True to launch the API instead of the webui")

 parser.add_argument("--ui-debug-mode", action='store_true', help="Don't load model to quickly launch UI")

 parser.add_argument("--device-id", type=str, help="Select the default CUDA device to use (export CUDA_VISIBLE_DEVICES=0,1,etc might be needed before)", default=None)

 parser.add_argument("--administrator", action='store_true', help="Administrator rights", default=False)

-parser.add_argument("--cors-allow-origins", type=str, help="Allowed CORS origins", default=None)

+parser.add_argument("--cors-allow-origins", type=str, help="Allowed CORS origin(s) in the form of a comma-separated list (no spaces)", default=None)

+parser.add_argument("--cors-allow-origins-regex", type=str, help="Allowed CORS origin(s) in the form of a single regular expression", default=None)

 parser.add_argument("--tls-keyfile", type=str, help="Partially enables TLS, requires --tls-certfile to fully function", default=None)

 parser.add_argument("--tls-certfile", type=str, help="Partially enables TLS, requires --tls-keyfile to fully function", default=None)

 parser.add_argument("--server-name", type=str, help="Sets hostname of server", default=None)

diff --git a/webui.py b/webui.py
index f4f1d74d..066d94f7 100644
--- a/webui.py
+++ b/webui.py
@@ -107,8 +107,12 @@ def initialize():
 

 

 def setup_cors(app):

-    if cmd_opts.cors_allow_origins:

+    if cmd_opts.cors_allow_origins and cmd_opts.cors_allow_origins_regex:

+        app.add_middleware(CORSMiddleware, allow_origins=cmd_opts.cors_allow_origins.split(','), allow_origin_regex=cmd_opts.cors_allow_origins_regex, allow_methods=['*'])

+    elif cmd_opts.cors_allow_origins:

         app.add_middleware(CORSMiddleware, allow_origins=cmd_opts.cors_allow_origins.split(','), allow_methods=['*'])

+    elif cmd_opts.cors_allow_origins_regex:

+        app.add_middleware(CORSMiddleware, allow_origin_regex=cmd_opts.cors_allow_origins_regex, allow_methods=['*'])

 

 

 def create_api(app):