From 44d8068bab4e603fe4eae0321e61338ecb7fa89f Mon Sep 17 00:00:00 2001
From: henk717 <henk717@msn.com>
Date: Mon, 29 Nov 2021 18:11:14 +0100
Subject: [PATCH] Ngrok Support

Not recommended for home users due to DDoS risks, but might make Colab tunnels more reliable.
---
 aiserver.py                    | 12 ++++++++++--
 environments/finetuneanon.yml  |  3 ++-
 environments/huggingface.yml   |  3 ++-
 environments/rocm-finetune.yml |  1 +
 environments/rocm.yml          |  1 +
 5 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/aiserver.py b/aiserver.py
index 47fae0e0..28da1d21 100644
--- a/aiserver.py
+++ b/aiserver.py
@@ -316,6 +316,7 @@ def device_config(model):
 # Parsing Parameters
 parser = argparse.ArgumentParser(description="KoboldAI Server")
 parser.add_argument("--remote", action='store_true', help="Optimizes KoboldAI for Remote Play")
+parser.add_argument("--ngrok", action='store_true', help="Optimizes KoboldAI for Remote Play using Ngrok")
 parser.add_argument("--model", help="Specify the Model Type to skip the Menu")
 parser.add_argument("--path", help="Specify the Path for local models (For model NeoCustom or GPT2Custom)")
 parser.add_argument("--cpu", action='store_true', help="By default unattended launches are on the GPU use this option to force CPU usage.")
@@ -332,6 +333,9 @@ vars.model = args.model;
 if args.remote:
     vars.remote = True;
 
+if args.ngrok:
+    vars.remote = True;
+
 vars.smandelete = vars.remote == args.override_delete
 vars.smanrename = vars.remote == args.override_rename
 
@@ -2929,8 +2933,12 @@ if __name__ == "__main__":
     
     #socketio.run(app, host='0.0.0.0', port=5000)
     if(vars.remote):
-        from flask_cloudflared import _run_cloudflared
-        cloudflare = _run_cloudflared(5000)
+        if(args.ngrok):
+            from flask_ngrok import _run_ngrok
+            cloudflare = _run_ngrok()
+        else:
+           from flask_cloudflared import _run_cloudflared
+           cloudflare = _run_cloudflared(5000)
         with open('cloudflare.log', 'w') as cloudflarelog:
             cloudflarelog.write("KoboldAI has finished loading and is available in the following link : " + cloudflare)
             print(format(colors.GREEN) + "KoboldAI has finished loading and is available in the following link : " + cloudflare + format(colors.END))
diff --git a/environments/finetuneanon.yml b/environments/finetuneanon.yml
index aefbdc87..76ae366b 100644
--- a/environments/finetuneanon.yml
+++ b/environments/finetuneanon.yml
@@ -14,4 +14,5 @@ dependencies:
   - git
   - pip:
     - git+https://github.com/finetuneanon/transformers@gpt-neo-localattention3-rp-b
-    - flask-cloudflared
\ No newline at end of file
+    - flask-cloudflared
+    - flask-ngrok
\ No newline at end of file
diff --git a/environments/huggingface.yml b/environments/huggingface.yml
index 41031ba8..b67c8a4a 100644
--- a/environments/huggingface.yml
+++ b/environments/huggingface.yml
@@ -14,4 +14,5 @@ dependencies:
   - pip
   - git
   - pip:
-    - flask-cloudflared
\ No newline at end of file
+    - flask-cloudflared
+    - flask-ngrok
\ No newline at end of file
diff --git a/environments/rocm-finetune.yml b/environments/rocm-finetune.yml
index 477dd6aa..8e281f68 100644
--- a/environments/rocm-finetune.yml
+++ b/environments/rocm-finetune.yml
@@ -14,3 +14,4 @@ dependencies:
     - torchvision==0.11.1
     - flask-cloudflared
     - git+https://github.com/finetuneanon/transformers@gpt-neo-localattention3-rp-b
+    - flask-ngrok
\ No newline at end of file
diff --git a/environments/rocm.yml b/environments/rocm.yml
index 237f2f8f..bbcb810b 100644
--- a/environments/rocm.yml
+++ b/environments/rocm.yml
@@ -14,3 +14,4 @@ dependencies:
     - torchvision==0.11.1
     - flask-cloudflared
     - git+https://github.com/vfbd/transformers#tfs
+    - flask-ngrok
\ No newline at end of file