Mightys/Notebook_Scripts

1import atexit2import os3import platform4import re5import stat6import subprocess7import sys8import time9from pathlib import Path10from typing import List, Optional11 12import requests13 14VERSION = "0.2"15CURRENT_TUNNELS: List["Tunnel"] = []16 17machine = platform.machine()18if machine == "x86_64":19    machine = "amd64"20 21BINARY_REMOTE_NAME = f"frpc_{platform.system().lower()}_{machine.lower()}"22EXTENSION = ".exe" if os.name == "nt" else ""23BINARY_URL = f"https://cdn-media.huggingface.co/frpc-gradio-{VERSION}/{BINARY_REMOTE_NAME}{EXTENSION}"24 25BINARY_FILENAME = f"{BINARY_REMOTE_NAME}_v{VERSION}"26BINARY_FOLDER = Path(__file__).parent.absolute()27BINARY_PATH = f"{BINARY_FOLDER / BINARY_FILENAME}"28 29TUNNEL_TIMEOUT_SECONDS = 3030TUNNEL_ERROR_MESSAGE = (31    "Could not create share URL. "32    "Please check the appended log from frpc for more information:"33)34 35GRADIO_API_SERVER = "https://api.gradio.app/v2/tunnel-request"36GRADIO_SHARE_SERVER_ADDRESS = None37 38 39class Tunnel:40    def __init__(self, remote_host, remote_port, local_host, local_port, share_token):41        self.proc = None42        self.url = None43        self.remote_host = remote_host44        self.remote_port = remote_port45        self.local_host = local_host46        self.local_port = local_port47        self.share_token = share_token48 49    @staticmethod50    def download_binary():51        if not Path(BINARY_PATH).exists():52            resp = requests.get(BINARY_URL)53 54            if resp.status_code == 403:55                raise OSError(56                    f"Cannot set up a share link as this platform is incompatible. Please "57                    f"create a GitHub issue with information about your platform: {platform.uname()}"58                )59 60            resp.raise_for_status()61 62            # Save file data to local copy63            with open(BINARY_PATH, "wb") as file:64                file.write(resp.content)65            st = os.stat(BINARY_PATH)66            os.chmod(BINARY_PATH, st.st_mode | stat.S_IEXEC)67 68    def start_tunnel(self) -> str:69        self.download_binary()70        self.url = self._start_tunnel(BINARY_PATH)71        return self.url72 73    def kill(self):74        if self.proc is not None:75            print(f"Killing tunnel {self.local_host}:{self.local_port} <> {self.url}")76            self.proc.terminate()77            self.proc = None78 79    def _start_tunnel(self, binary: str) -> str:80        CURRENT_TUNNELS.append(self)81        command = [82            binary,83            "http",84            "-n",85            self.share_token,86            "-l",87            str(self.local_port),88            "-i",89            self.local_host,90            "--uc",91            "--sd",92            "random",93            "--ue",94            "--server_addr",95            f"{self.remote_host}:{self.remote_port}",96            "--disable_log_color",97        ]98        self.proc = subprocess.Popen(99            command, stdout=subprocess.PIPE, stderr=subprocess.PIPE100        )101        atexit.register(self.kill)102        return self._read_url_from_tunnel_stream()103 104    def _read_url_from_tunnel_stream(self) -> str:105        start_timestamp = time.time()106 107        log = []108        url = ""109 110        def _raise_tunnel_error():111            log_text = "\n".join(log)112            print(log_text, file=sys.stderr)113            raise ValueError(f"{TUNNEL_ERROR_MESSAGE}\n{log_text}")114 115        while url == "":116            # check for timeout and log117            if time.time() - start_timestamp >= TUNNEL_TIMEOUT_SECONDS:118                _raise_tunnel_error()119 120            assert self.proc is not None121            if self.proc.stdout is None:122                continue123 124            line = self.proc.stdout.readline()125            line = line.decode("utf-8")126 127            if line == "":128                continue129 130            log.append(line.strip())131 132            if "start proxy success" in line:133                result = re.search("start proxy success: (.+)\n", line)134                if result is None:135                    _raise_tunnel_error()136                else:137                    url = result.group(1)138            elif "login to server failed" in line:139                _raise_tunnel_error()140 141        return url142 143 144def setup_tunnel(145    local_host: str,146    local_port: int,147    share_token: str,148    share_server_address: Optional[str],149) -> str:150    share_server_address = (151        GRADIO_SHARE_SERVER_ADDRESS152        if share_server_address is None153        else share_server_address154    )155    if share_server_address is None:156        response = requests.get(GRADIO_API_SERVER)157        if not (response and response.status_code == 200):158            raise RuntimeError("Could not get share link from Gradio API Server.")159        payload = response.json()[0]160        remote_host, remote_port = payload["host"], int(payload["port"])161    else:162        remote_host, remote_port = share_server_address.split(":")163        remote_port = int(remote_port)164    try:165        tunnel = Tunnel(remote_host, remote_port, local_host, local_port, share_token)166        address = tunnel.start_tunnel()167        return address168    except Exception as e:169        raise RuntimeError(str(e)) from e170 171 172def main():173    import argparse174    import secrets175    import time176 177    parser = argparse.ArgumentParser(description="Set up a tunnel.")178    parser.add_argument(179        "-p",180        "--port",181        type=int,182        default=8080,183        help="the port number to use for the tunnel.",184    )185    parser.add_argument(186        "port_positional",187        type=int,188        nargs="?",189        help="the port number to use for the tunnel.",190    )191    parser.add_argument(192        "--sd", "--subdomain", type=str, help="the subdomain to use for the tunnel."193    )194    args = parser.parse_args()195 196    if args.port_positional is not None:197        args.port = args.port_positional198 199    address = setup_tunnel(200        "127.0.0.1",201        args.port,202        secrets.token_urlsafe(32) if args.sd is None else args.sd,203        None,204    )205 206    print(address, flush=True)207    time.sleep(3600 * 24 * 3)208 209if __name__ == "__main__":210    main()211