Use ngrok-py library

3 files changed, 17 insertions, 25 deletions

diff --git a/modules/cmd_args.py b/modules/cmd_args.py
index d906a571..bf18b7b7 100644
--- a/modules/cmd_args.py
+++ b/modules/cmd_args.py
@@ -1,4 +1,5 @@
 import argparse

+import json

 import os

 from modules.paths_internal import models_path, script_path, data_path, extensions_dir, extensions_builtin_dir, sd_default_config, sd_model_file

 

@@ -39,7 +40,7 @@ parser.add_argument("--precision", type=str, help="evaluate at this precision",
 parser.add_argument("--upcast-sampling", action='store_true', help="upcast sampling. No effect with --no-half. Usually produces similar results to --no-half with better performance while using less memory.")

 parser.add_argument("--share", action='store_true', help="use share=True for gradio and make the UI accessible through their site")

 parser.add_argument("--ngrok", type=str, help="ngrok authtoken, alternative to gradio --share", default=None)

-parser.add_argument("--ngrok-region", type=str, help="The region in which ngrok should start.", default="us")

+parser.add_argument("--ngrok-options", type=json.loads, help='The options to pass to ngrok in JSON format, e.g.: \'{"authtoken_from_env":true, "basic_auth":"user:password", "oauth_provider":"google", "oauth_allow_emails":"user@asdf.com"}\'', default=dict())

 parser.add_argument("--enable-insecure-extension-access", action='store_true', help="enable extensions tab regardless of other options")

 parser.add_argument("--codeformer-models-path", type=str, help="Path to directory with codeformer model file(s).", default=os.path.join(models_path, 'Codeformer'))

 parser.add_argument("--gfpgan-models-path", type=str, help="Path to directory with GFPGAN model file(s).", default=os.path.join(models_path, 'GFPGAN'))

diff --git a/modules/ngrok.py b/modules/ngrok.py
index 7a7b4b26..caa352d1 100644
--- a/modules/ngrok.py
+++ b/modules/ngrok.py
@@ -1,6 +1,7 @@
-from pyngrok import ngrok, conf, exception
+import ngrok
 
-def connect(token, port, region):
+# Connect to ngrok for ingress
+def connect(token, port, options):
     account = None
     if token is None:
         token = 'None'
@@ -10,28 +11,18 @@ def connect(token, port, region):
             token, username, password = token.split(':', 2)
             account = f"{username}:{password}"
 
-    config = conf.PyngrokConfig(
-        auth_token=token, region=region
-    )
-    
-    # Guard for existing tunnels
-    existing = ngrok.get_tunnels(pyngrok_config=config)
-    if existing:
-        for established in existing:
-            # Extra configuration in the case that the user is also using ngrok for other tunnels
-            if established.config['addr'][-4:] == str(port):
-                public_url = existing[0].public_url
-                print(f'ngrok has already been connected to localhost:{port}! URL: {public_url}\n'
-                    'You can use this link after the launch is complete.')
-                return
-    
+    # For all options see: https://github.com/ngrok/ngrok-py/blob/main/examples/ngrok-connect-full.py
+    if not options.get('authtoken_from_env'):
+        options['authtoken'] = token
+    if account:
+        options['basic_auth'] = account
+    if not options.get('session_metadata'):
+        options['session_metadata'] = 'stable-diffusion-webui'
+
     try:
-        if account is None:
-            public_url = ngrok.connect(port, pyngrok_config=config, bind_tls=True).public_url
-        else:
-            public_url = ngrok.connect(port, pyngrok_config=config, bind_tls=True, auth=account).public_url
-    except exception.PyngrokNgrokError:
-        print(f'Invalid ngrok authtoken, ngrok connection aborted.\n'
+        public_url = ngrok.connect(f"127.0.0.1:{port}", **options).url()
+    except Exception as e:
+        print(f'Invalid ngrok authtoken? ngrok connection aborted due to: {e}\n'
               f'Your token: {token}, get the right one on https://dashboard.ngrok.com/get-started/your-authtoken')
     else:
         print(f'ngrok connected to localhost:{port}! URL: {public_url}\n'
diff --git a/modules/ui.py b/modules/ui.py
index f07bcc41..5f5405f0 100644
--- a/modules/ui.py
+++ b/modules/ui.py
@@ -59,7 +59,7 @@ if cmd_opts.ngrok is not None:
     ngrok.connect(

         cmd_opts.ngrok,

         cmd_opts.port if cmd_opts.port is not None else 7860,

-        cmd_opts.ngrok_region

+        cmd_opts.ngrok_options

         )