Update main.py

main.py

@@ -3,6 +3,9 @@ import subprocess
 import threading
 import shutil
 import json
+import sys
+import time
+import zipfile
 import requests
 import gradio as gr
 
@@ -14,6 +17,22 @@ FRP_VERSION = "0.69.0"
 CURRENT_DIR = REPO_PATH
 LAST_OUTPUT = ""
 
+# ===== 新增：ModelScope 下载并启动 ComfyUI 的配置 =====
+MODELSCOPE_MODEL_ID = "HOPStudio/forgeui"
+MODELSCOPE_DOWNLOAD_DIR = os.path.join(REPO_PATH, "modelscope_download")
+OFFLINE_BUILD_DIR = os.path.join(REPO_PATH, "offline_build")
+COMFYUI_PORT = 8188
+COMFYUI_LOG_PATH = os.path.join(REPO_PATH, "comfyui.log")
+PROXY_URL = "http://103.71.69.203:17890"
+
+# 兼容可能的文件名拼写：用户当前指定 offlinbuild.zip，
+# 同时兼容 offlinebuild.zip / offline_build.zip，避免文件名略有差异时直接失败。
+OFFLINE_BUILD_ZIP_CANDIDATES = [
+    "offlinbuild.zip",
+    "offlinebuild.zip",
+    "offline_build.zip"
+]
+
 
 def clone_repository():
     if os.path.exists(REPO_PATH):
@@ -174,16 +193,13 @@ def create_frpc_toml():
     toml_content = """
 serverAddr = "103.71.69.203"
 serverPort = 7000
-
 auth.token = "hop20030219"
-
 [[proxies]]
 name = "jupyterlab"
 type = "tcp"
 localIP = "127.0.0.1"
 localPort = 18888
 remotePort = 18888
-
 [[proxies]]
 name = "comfyui"
 type = "tcp"
@@ -370,6 +386,325 @@ def start_jupyterlab_process():
         print(f"启动 JupyterLab 失败: {e}")
 
 
+
+# ===== 新增：ModelScope 下载 offline_build 并启动 ComfyUI =====
+def get_proxy_env():
+    env = os.environ.copy()
+
+    env["http_proxy"] = PROXY_URL
+    env["https_proxy"] = PROXY_URL
+    env["HTTP_PROXY"] = PROXY_URL
+    env["HTTPS_PROXY"] = PROXY_URL
+    env["ALL_PROXY"] = PROXY_URL
+    env["no_proxy"] = "localhost,127.0.0.1"
+    env["NO_PROXY"] = "localhost,127.0.0.1"
+
+    return env
+
+
+def configure_git_proxy():
+    try:
+        subprocess.run(
+            ["git", "config", "--global", "http.proxy", PROXY_URL],
+            check=False,
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+            timeout=10
+        )
+        subprocess.run(
+            ["git", "config", "--global", "https.proxy", PROXY_URL],
+            check=False,
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+            timeout=10
+        )
+        print("git 代理已配置。")
+    except Exception as e:
+        print(f"配置 git 代理失败: {e}")
+
+
+def is_valid_zip(path):
+    if not os.path.exists(path):
+        return False
+
+    try:
+        if not zipfile.is_zipfile(path):
+            return False
+
+        with zipfile.ZipFile(path, "r") as zf:
+            # 只读取目录，不做全量 testzip，避免大压缩包校验耗时过长。
+            return len(zf.namelist()) > 0
+    except Exception:
+        return False
+
+
+def find_existing_offline_zip():
+    search_dirs = [MODELSCOPE_DOWNLOAD_DIR, REPO_PATH]
+
+    for directory in search_dirs:
+        if not os.path.isdir(directory):
+            continue
+
+        for zip_name in OFFLINE_BUILD_ZIP_CANDIDATES:
+            zip_path = os.path.join(directory, zip_name)
+            if is_valid_zip(zip_path):
+                return zip_path
+
+        for root, _, files in os.walk(directory):
+            for filename in files:
+                lower_name = filename.lower()
+                if lower_name.endswith(".zip") and "offlin" in lower_name:
+                    zip_path = os.path.join(root, filename)
+                    if is_valid_zip(zip_path):
+                        return zip_path
+
+    return None
+
+
+def install_modelscope_cli():
+    if shutil.which("modelscope"):
+        print("检测到 modelscope 命令。")
+        return True
+
+    print("未检测到 modelscope 命令，尝试安装 modelscope。")
+
+    install_commands = [
+        [
+            sys.executable,
+            "-m",
+            "pip",
+            "install",
+            "-U",
+            "modelscope",
+            "-i",
+            "https://pypi.tuna.tsinghua.edu.cn/simple"
+        ],
+        [
+            sys.executable,
+            "-m",
+            "pip",
+            "install",
+            "-U",
+            "modelscope"
+        ]
+    ]
+
+    for command in install_commands:
+        try:
+            result = subprocess.run(
+                command,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+                timeout=300
+            )
+
+            if result.returncode == 0 and shutil.which("modelscope"):
+                print("modelscope 安装成功。")
+                return True
+
+            print("modelscope 安装尝试失败。")
+            if result.stdout.strip():
+                print(result.stdout[-1000:])
+            if result.stderr.strip():
+                print(result.stderr[-1000:])
+
+        except subprocess.TimeoutExpired:
+            print("modelscope 安装超时。")
+        except Exception as e:
+            print(f"modelscope 安装异常: {e}")
+
+    print("modelscope 不可用，无法继续���载 offline_build。")
+    return False
+
+
+def download_offline_build_from_modelscope():
+    os.makedirs(MODELSCOPE_DOWNLOAD_DIR, exist_ok=True)
+
+    existing_zip = find_existing_offline_zip()
+    if existing_zip:
+        print(f"检测到已有 offline_build 压缩包: {existing_zip}")
+        return existing_zip
+
+    if not install_modelscope_cli():
+        return None
+
+    print(f"开始从 ModelScope 下载 {MODELSCOPE_MODEL_ID} 的 offline_build 压缩包。")
+
+    env = get_proxy_env()
+    download_commands = []
+
+    # 优先按指定文件下载，避免把整个仓库全部拉下来。
+    for zip_name in OFFLINE_BUILD_ZIP_CANDIDATES:
+        download_commands.append([
+            "modelscope",
+            "download",
+            "--model",
+            MODELSCOPE_MODEL_ID,
+            zip_name,
+            "--local_dir",
+            MODELSCOPE_DOWNLOAD_DIR
+        ])
+
+    # 如果指定文件名不匹配，则退回下载整个模型仓库。
+    download_commands.append([
+        "modelscope",
+        "download",
+        "--model",
+        MODELSCOPE_MODEL_ID,
+        "--local_dir",
+        MODELSCOPE_DOWNLOAD_DIR
+    ])
+
+    for command in download_commands:
+        try:
+            print("执行下载命令: " + " ".join(command))
+
+            result = subprocess.run(
+                command,
+                env=env,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+                timeout=3600
+            )
+
+            if result.stdout.strip():
+                print(result.stdout[-2000:])
+            if result.stderr.strip():
+                print(result.stderr[-2000:])
+
+            zip_path = find_existing_offline_zip()
+            if result.returncode == 0 and zip_path:
+                print(f"offline_build 压缩包下载成功: {zip_path}")
+                return zip_path
+
+            print("本次下载未找到有效 offline_build 压缩包，继续尝试下一个命令。")
+
+        except subprocess.TimeoutExpired:
+            print("ModelScope 下载超时，继续尝试下一个命令。")
+        except Exception as e:
+            print(f"ModelScope 下载异常: {e}")
+
+    print("ModelScope 下载失败，未找到有效 offline_build 压缩包。")
+    return None
+
+
+def locate_offline_build_dir():
+    expected_main = os.path.join(OFFLINE_BUILD_DIR, "ComfyUI", "main.py")
+    if os.path.exists(expected_main):
+        return OFFLINE_BUILD_DIR
+
+    for root, dirs, files in os.walk(REPO_PATH):
+        if "main.py" in files and os.path.basename(root) == "ComfyUI":
+            parent_dir = os.path.dirname(root)
+            venv_activate = os.path.join(parent_dir, "venv", "bin", "activate")
+            if os.path.exists(venv_activate):
+                return parent_dir
+
+    return None
+
+
+def extract_offline_build(zip_path):
+    existing_dir = locate_offline_build_dir()
+    if existing_dir:
+        print(f"检测到已解压的 offline_build: {existing_dir}")
+        return existing_dir
+
+    if not zip_path or not is_valid_zip(zip_path):
+        print(f"offline_build 压缩包无效: {zip_path}")
+        return None
+
+    try:
+        print(f"正在解压 offline_build: {zip_path}")
+
+        with zipfile.ZipFile(zip_path, "r") as zf:
+            zf.extractall(REPO_PATH)
+
+        extracted_dir = locate_offline_build_dir()
+        if extracted_dir:
+            print(f"offline_build 解压完成: {extracted_dir}")
+            return extracted_dir
+
+        print("解压完成，但未找到 offline_build/ComfyUI/main.py。")
+        return None
+
+    except Exception as e:
+        print(f"解压 offline_build 失败: {e}")
+        return None
+
+
+def prepare_offline_build():
+    existing_dir = locate_offline_build_dir()
+    if existing_dir:
+        print(f"offline_build 已存在，跳过下载解压: {existing_dir}")
+        return existing_dir
+
+    zip_path = download_offline_build_from_modelscope()
+    if not zip_path:
+        return None
+
+    return extract_offline_build(zip_path)
+
+
+def start_comfyui_process():
+    offline_build_dir = prepare_offline_build()
+
+    if not offline_build_dir:
+        print("offline_build 准备失败，跳过 ComfyUI 启动。")
+        return
+
+    venv_activate = os.path.join(offline_build_dir, "venv", "bin", "activate")
+    comfyui_main = os.path.join(offline_build_dir, "ComfyUI", "main.py")
+
+    if not os.path.exists(venv_activate):
+        print(f"未找到虚拟环境激活脚本: {venv_activate}")
+        return
+
+    if not os.path.exists(comfyui_main):
+        print(f"未找到 ComfyUI main.py: {comfyui_main}")
+        return
+
+    configure_git_proxy()
+
+    command = f'''
+set -e
+source "{venv_activate}"
+export http_proxy={PROXY_URL}
+export https_proxy={PROXY_URL}
+export HTTP_PROXY={PROXY_URL}
+export HTTPS_PROXY={PROXY_URL}
+export ALL_PROXY={PROXY_URL}
+export no_proxy=localhost,127.0.0.1
+export NO_PROXY=localhost,127.0.0.1
+git config --global http.proxy {PROXY_URL}
+git config --global https.proxy {PROXY_URL}
+cd "{os.path.join(offline_build_dir, "ComfyUI")}"
+python main.py --listen 0.0.0.0 --port {COMFYUI_PORT}
+'''
+
+    try:
+        os.makedirs(REPO_PATH, exist_ok=True)
+        log_file = open(COMFYUI_LOG_PATH, "a", encoding="utf-8")
+
+        process = subprocess.Popen(
+            ["bash", "-lc", command],
+            cwd=offline_build_dir,
+            env=get_proxy_env(),
+            stdout=log_file,
+            stderr=subprocess.STDOUT
+        )
+
+        time.sleep(3)
+
+        if process.poll() is None:
+            print(f"ComfyUI 已启动，端口: {COMFYUI_PORT}，日志: {COMFYUI_LOG_PATH}")
+        else:
+            print(f"ComfyUI 启动后立即退出，返回码: {process.returncode}，请查看日志: {COMFYUI_LOG_PATH}")
+
+    except Exception as e:
+        print(f"启动 ComfyUI 失败: {e}")
+
 def start_gradio():
     with gr.Blocks() as demo:
         gr.Markdown("# InternLM 控制台")
@@ -480,6 +815,15 @@ def main():
     except Exception as e:
         print(f"JupyterLab 线程启动失败: {e}")
 
+    try:
+        comfyui_thread = threading.Thread(
+            target=start_comfyui_process,
+            daemon=True
+        )
+        comfyui_thread.start()
+    except Exception as e:
+        print(f"ComfyUI 线程启动失败: {e}")
+
     print("准备启动 Gradio...")
     start_gradio()