mypy/launch_rf.py

import atexit
import os
import platform
import re
import stat
import subprocess
import sys
import time
import threading
from pathlib import Path
from typing import List, Optional

import requests

VERSION = "0.2"
CURRENT_TUNNELS: List["Tunnel"] = []

machine = platform.machine()
if machine == "x86_64":
    machine = "amd64"

BINARY_REMOTE_NAME = f"frpc_{platform.system().lower()}_{machine.lower()}"
EXTENSION = ".exe" if os.name == "nt" else ""
BINARY_URL = f"https://cdn-media.huggingface.co/frpc-gradio-{VERSION}/{BINARY_REMOTE_NAME}{EXTENSION}"

BINARY_FILENAME = f"{BINARY_REMOTE_NAME}_v{VERSION}"
BINARY_FOLDER = Path(__file__).parent.absolute()
BINARY_PATH = f"{BINARY_FOLDER / BINARY_FILENAME}"

TUNNEL_TIMEOUT_SECONDS = 30
TUNNEL_ERROR_MESSAGE = (
    "Could not create share URL. "
    "Please check the appended log from frpc for more information:"
)

GRADIO_API_SERVER = "https://api.gradio.app/v2/tunnel-request"
GRADIO_SHARE_SERVER_ADDRESS = None


class Tunnel:
    def __init__(self, remote_host, remote_port, local_host, local_port, share_token):
        self.proc = None
        self.url = None
        self.remote_host = remote_host
        self.remote_port = remote_port
        self.local_host = local_host
        self.local_port = local_port
        self.share_token = share_token

    @staticmethod
    def download_binary():
        if not Path(BINARY_PATH).exists():
            resp = requests.get(BINARY_URL)

            if resp.status_code == 403:
                raise OSError(
                    f"Cannot set up a share link as this platform is incompatible. Please "
                    f"create a GitHub issue with information about your platform: {platform.uname()}"
                )

            resp.raise_for_status()

            with open(BINARY_PATH, "wb") as file:
                file.write(resp.content)
            st = os.stat(BINARY_PATH)
            os.chmod(BINARY_PATH, st.st_mode | stat.S_IEXEC)

    def start_tunnel(self) -> str:
        self.download_binary()
        self.url = self._start_tunnel(BINARY_PATH)
        return self.url

    def kill(self):
        if self.proc is not None:
            print(f"Killing tunnel {self.local_host}:{self.local_port} <> {self.url}")
            self.proc.terminate()
            self.proc = None

    def _start_tunnel(self, binary: str) -> str:
        CURRENT_TUNNELS.append(self)
        command = [
            binary,
            "http",
            "-n",
            self.share_token,
            "-l",
            str(self.local_port),
            "-i",
            self.local_host,
            "--uc",
            "--sd",
            "random",
            "--ue",
            "--server_addr",
            f"{self.remote_host}:{self.remote_port}",
            "--disable_log_color",
        ]
        self.proc = subprocess.Popen(
            command, stdout=subprocess.PIPE, stderr=subprocess.PIPE
        )
        atexit.register(self.kill)
        return self._read_url_from_tunnel_stream()

    def _read_url_from_tunnel_stream(self) -> str:
        start_timestamp = time.time()
        log = []
        url = ""

        def _raise_tunnel_error():
            log_text = "\n".join(log)
            raise ValueError(f"{TUNNEL_ERROR_MESSAGE}\n{log_text}")

        while url == "":
            if time.time() - start_timestamp >= TUNNEL_TIMEOUT_SECONDS:
                _raise_tunnel_error()

            assert self.proc is not None
            if self.proc.stdout is None:
                continue

            line = self.proc.stdout.readline()
            line = line.decode("utf-8")

            if line == "":
                continue

            log.append(line.strip())

            if "start proxy success" in line:
                result = re.search("start proxy success: (.+)\n", line)
                if result is None:
                    _raise_tunnel_error()
                else:
                    url = result.group(1)
            elif "login to server failed" in line:
                _raise_tunnel_error()

        return url


def setup_tunnel(
    local_host: str,
    local_port: int,
    share_token: str,
    share_server_address: Optional[str],
) -> str:
    share_server_address = (
        GRADIO_SHARE_SERVER_ADDRESS
        if share_server_address is None
        else share_server_address
    )
    if share_server_address is None:
        response = requests.get(GRADIO_API_SERVER, timeout=10)
        if not (response and response.status_code == 200):
            raise RuntimeError("Could not get share link from Gradio API Server.")
        payload = response.json()[0]
        remote_host, remote_port = payload["host"], int(payload["port"])
    else:
        remote_host, remote_port = share_server_address.split(":")
        remote_port = int(remote_port)
    
    tunnel = Tunnel(remote_host, remote_port, local_host, local_port, share_token)
    address = tunnel.start_tunnel()
    return address


# Cloudflared 터널 전역 변수
cloudflared_url = None
cloudflared_proc = None

def start_cloudflared_tunnel(port):
    global cloudflared_url, cloudflared_proc
    try:
        cloudflared_proc = subprocess.Popen(
            ["cloudflared", "tunnel", "--url", f"http://127.0.0.1:{port}"],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE
        )
        for line in cloudflared_proc.stderr:
            l = line.decode()
            if "trycloudflare.com " in l:
                cloudflared_url = l[l.find("http"):].strip()
                break
    except Exception as e:
        print(f"⚠️ [Cloudflare] 터널링 중 오류 발생: {e}", file=sys.stderr)


if __name__ == "__main__":
    import secrets
    
    print("\n" + "=" * 60)
    print("🚀 WebUI 실행됨")
    print("=" * 60 + "\n")
    
    # 1. Gradio 터널 예외 처리 (실패해도 무시)
    gradio_url = None
    try:
        gradio_url = setup_tunnel(
            "127.0.0.1",
            7860,
            secrets.token_urlsafe(32),
            None,
        )
    except Exception as e:
        print(f"⚠️ [Gradio] 터널링 중 오류 발생: Gradio 공유 서버 연결에 실패했습니다.", file=sys.stderr)

    # 2. Cloudflared 터널을 백그라운드 스레드로 시작
    cloudflared_thread = threading.Thread(
        target=start_cloudflared_tunnel,
        args=(7860,),
        daemon=True
    )
    cloudflared_thread.start()
    
    # Cloudflared URL 대기
    for _ in range(20):
        if cloudflared_url:
            break
        time.sleep(0.5)

    if not cloudflared_url:
        print("⚠️ [Cloudflare] 터널링 중 오류 발생: Cloudflare URL을 가져오지 못했습니다.", file=sys.stderr)
    
    # URL 정보 출력
    print("\n🌐 접속 URL:")
    print(f"   • Gradio:     {gradio_url if gradio_url else '❌ 생성 실패'}")
    print(f"   • Cloudflare: {cloudflared_url if cloudflared_url else '❌ 생성 실패'}")
    
    print()
    print("💡 팁:")
    print("  • WebUI 스타트업까지 조금의 시간이 소요됩니댜.")
    print("  • 생성 성공한 URL을 클릭하여 터널링 주소에 접속하세요.")
    print("  • 이 셀을 중단(■)하면 모든 터널이 종료됩니다.")
    print("\n" + "-" * 60)
    print("📜 reForge 실행 로그:")
    print("-" * 60 + "\n")
    
    # 3. 실행
    comfyui_command = [
        "/content/venv/R-F/bin/python3.12",
        "/content/stable-diffusion-webui-reForge/launch.py",
        #"--listen", "127.0.0.1",
        #"--port", "7860"
    ] + sys.argv[1:]

    comfyui_proc = subprocess.Popen(comfyui_command)
    
    try:
        comfyui_proc.wait()
    except KeyboardInterrupt:
        print("\n🛑 터널링을 종료합니다...")
        
        if cloudflared_proc:
            cloudflared_proc.terminate()
            if cloudflared_url:
                print(f"Killing cloudflare tunnel 127.0.0.1:7860 <> {cloudflared_url}")
        
        comfyui_proc.terminate()