| import atexit |
| import re |
| import shlex |
| import subprocess |
| from pathlib import Path |
| from tempfile import TemporaryDirectory |
| from typing import Union |
| from gradio import strings |
| import os |
|
|
| from modules.shared import cmd_opts |
|
|
| LOCALHOST_RUN = "localhost.run" |
| REMOTE_MOE = "remote.moe" |
| localhostrun_pattern = re.compile(r"(?P<url>https?://\S+\.lhr\.life)") |
| remotemoe_pattern = re.compile(r"(?P<url>https?://\S+\.remote\.moe)") |
|
|
|
|
| def gen_key(path: Union[str, Path]) -> None: |
| path = Path(path) |
| arg_string = f'ssh-keygen -t rsa -b 4096 -N "" -q -f {path.as_posix()}' |
| args = shlex.split(arg_string) |
| subprocess.run(args, check=True) |
| path.chmod(0o600) |
|
|
|
|
| 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) |
| |
| 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:127.0.0.1:{port} -o StrictHostKeyChecking=no -i {ssh_path.as_posix()} {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}") |
|
|
| |
| os.environ['webui_url'] = tunnel_url |
| colab_url = os.getenv('colab_url') |
| strings.en["SHARE_LINK_MESSAGE"] = f"Running on public URL (recommended): {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) |
|
|
