From 0d31f20cbd556ea4ba3d8ad9254bcce71c32088c Mon Sep 17 00:00:00 2001
From: bobzilladev <bob@ngrok.com>
Date: Tue, 16 May 2023 13:15:30 -0400
Subject: [PATCH] Use ngrok-py library

---
 launch.py           |  4 ++--
 modules/cmd_args.py |  3 ++-
 modules/ngrok.py    | 37 ++++++++++++++-----------------------
 modules/ui.py       |  2 +-
 4 files changed, 19 insertions(+), 27 deletions(-)

diff --git a/launch.py b/launch.py
index cfc0cffa..871d3ea8 100644
--- a/launch.py
+++ b/launch.py
@@ -294,8 +294,8 @@ def prepare_environment():
         elif platform.system() == "Linux":
             run_pip(f"install {xformers_package}", "xformers")
 
-    if not is_installed("pyngrok") and args.ngrok:
-        run_pip("install pyngrok", "ngrok")
+    if not is_installed("ngrok") and args.ngrok:
+        run_pip("install ngrok", "ngrok")
 
     os.makedirs(os.path.join(script_path, dir_repos), exist_ok=True)
 
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
         )