Upload launch_neo.py

mypy/launch_neo.py

@@ -0,0 +1,258 @@
+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("📜 Forge neo 실행 로그:")
+    print("-" * 60 + "\n")
+    
+    # 3. 실행
+    comfyui_command = [
+        "/content/venv/NEO/bin/python3.13",
+        "/content/sd-webui-forge-neo/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()