feat: add ssh tunnels

.gitignore

@@ -171,3 +171,6 @@ poetry.toml
 .ruff_cache/
 
 # End of https://www.toptal.com/developers/gitignore/api/python
+
+id_rsa
+id_rsa.pub

README.md

@@ -4,6 +4,14 @@ Tunneling extension for [AUTOMATIC1111/stable-diffusion-webui](https://github.co
 
 ## Usage
 
-### cloudflared
+### [cloudflared](https://try.cloudflare.com/)
 
 add `--cloudflared` to commandline options.
+
+### [localhost.run](https://localhost.run/)
+
+add `--localhostrun` to commandline options.
+
+### [remote.moe](https://github.com/fasmide/remotemoe)
+
+add `--remotemoe` to commandline options.

preload.py

@@ -7,3 +7,15 @@ def preload(parser: argparse.ArgumentParser):
         action="store_true",
         help="use trycloudflare, alternative to gradio --share",
     )
+
+    parser.add_argument(
+        "--localhostrun",
+        action="store_true",
+        help="use localhost.run, alternative to gradio --share",
+    )
+
+    parser.add_argument(
+        "--remotemoe",
+        action="store_true",
+        help="use remote.moe, alternative to gradio --share",
+    )

scripts/_util.py

@@ -0,0 +1,11 @@
+from __future__ import annotations
+
+import shlex
+import subprocess
+from pathlib import Path
+
+
+def gen_key(path: str | Path) -> None:
+    arg_string = f'ssh-keygen -t rsa -b 4096 -N "" -f {path}'
+    args = shlex.split(arg_string)
+    subprocess.run(args, check=True)

scripts/ssh_tunnel.py

@@ -0,0 +1,68 @@
+import atexit
+import re
+import shlex
+import subprocess
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+from modules.shared import cmd_opts
+
+from ._util import gen_key
+
+LOCALHOST_RUN = "localhost.run"
+localhostrun_pattern = re.compile(r"(?P<url>https?://\S+\.lhr\.life)")
+remotemoe_pattern = re.compile(r"(?P<url>https?://\S+\.remote\.moe)")
+
+
+def ssh_tunnel(host: str = LOCALHOST_RUN) -> None:
+    ssh_name = "id_rsa"
+    ssh_path = Path(__file__).parent.parent / ssh_name
+
+    tmp = None
+    if not ssh_path.exists():
+        try:
+            gen_key(ssh_path)
+        # write permission error or etc
+        except subprocess.CalledProcessError:
+            tmp = TemporaryDirectory()
+            ssh_path = Path(tmp.name) / ssh_name
+            gen_key(ssh_path)
+
+    port = cmd_opts.port if cmd_opts.port else 7860
+    arg_string = (
+        f"ssh -R 80:localhost:{port} -o StrictHostKeyChecking=no -i {ssh_path} {host}"
+    )
+    args = shlex.split(arg_string)
+    tunnel = subprocess.Popen(
+        args, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, encoding="utf-8"
+    )
+
+    atexit.register(tunnel.terminate)
+    if tmp is not None:
+        atexit.register(tmp.cleanup)
+
+    tunnel_url = ""
+    lines = 27 if host == LOCALHOST_RUN else 5
+    pattern = localhostrun_pattern if host == LOCALHOST_RUN else remotemoe_pattern
+    for _ in range(lines):
+        line = tunnel.stdout.readline()
+        if line.startswith("Warning"):
+            print(line, end="")
+
+        url_match = pattern.search(line)
+        if url_match:
+            tunnel_url = url_match.group("url")
+            break
+    else:
+        raise RuntimeError(f"Failed to run {host}")
+
+    print(f" * Running on {tunnel_url}")
+
+
+if cmd_opts.localhostrun:
+    print("localhost.run detected, trying to connect...")
+    ssh_tunnel(LOCALHOST_RUN)
+
+if cmd_opts.remotemoe:
+    print("remote.moe detected, trying to connect...")
+    ssh_tunnel("remote.moe")