Update main.py

main.py

@@ -6,6 +6,7 @@ import json
 import sys
 import time
 import zipfile
+import re
 import requests
 import gradio as gr
 
@@ -17,13 +18,17 @@ FRP_VERSION = "0.69.0"
 CURRENT_DIR = REPO_PATH
 LAST_OUTPUT = ""
 
-# ===== 新增：ModelScope 下载并启动 ComfyUI 的配置 =====
+# ===== 新增：ModelScope 整库克隆/下载并启动 ComfyUI 的配置 =====
 MODELSCOPE_MODEL_ID = "HOPStudio/forgeui"
-MODELSCOPE_DOWNLOAD_DIR = os.path.join(REPO_PATH, "modelscope_download")
+# forgeui 整库和 offline_build 下载包都很大，不放在 REPO_PATH 里，避免 clone_repository() 删除 /root/modelscope 时把断点文件删掉。
+OFFLINE_BUILD_CACHE_DIR = "/root/offline_build_cache"
+MODELSCOPE_DOWNLOAD_DIR = os.path.join(OFFLINE_BUILD_CACHE_DIR, "modelscope_download")
+MODELSCOPE_CACHE_DIR = os.path.join(OFFLINE_BUILD_CACHE_DIR, "modelscope_cache")
 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"
+MODELSCOPE_DOWNLOAD_TIMEOUT = 21600
 
 # 兼容可能的文件名拼写：用户当前指定 offlinbuild.zip，
 # 同时兼容 offlinebuild.zip / offline_build.zip，避免文件名略有差异时直接失败。
@@ -387,7 +392,7 @@ def start_jupyterlab_process():
 
 
 
-# ===== 新增：ModelScope 下载 offline_build 并启动 ComfyUI =====
+# ===== 新增：ModelScope 整库克隆/下载 offline_build 并启动 ComfyUI =====
 def get_proxy_env():
     env = os.environ.copy()
 
@@ -402,6 +407,43 @@ def get_proxy_env():
     return env
 
 
+def get_modelscope_download_env(use_proxy=False):
+    env = os.environ.copy()
+
+    # 固定 ModelScope 缓存目录，避免默认缓存散落到 ~/.cache，且方便断点续传。
+    os.makedirs(MODELSCOPE_CACHE_DIR, exist_ok=True)
+    env["MODELSCOPE_CACHE"] = MODELSCOPE_CACHE_DIR
+    env["MODELSCOPE_ENDPOINT"] = "https://www.modelscope.cn"
+    env["MODELSCOPE_SDK_ENDPOINT"] = "https://www.modelscope.cn"
+
+    # ModelScope 下载大文件时会跳到阿里云 OSS。这里默认直连，避免代理把 OSS 内网地址转坏。
+    proxy_keys = [
+        "http_proxy",
+        "https_proxy",
+        "HTTP_PROXY",
+        "HTTPS_PROXY",
+        "ALL_PROXY",
+        "all_proxy"
+    ]
+
+    if use_proxy:
+        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["all_proxy"] = PROXY_URL
+    else:
+        for key in proxy_keys:
+            env.pop(key, None)
+
+    no_proxy_value = "localhost,127.0.0.1,modelscope.cn,www.modelscope.cn,.modelscope.cn,aliyuncs.com,.aliyuncs.com"
+    env["no_proxy"] = no_proxy_value
+    env["NO_PROXY"] = no_proxy_value
+
+    return env
+
+
 def configure_git_proxy():
     try:
         subprocess.run(
@@ -439,7 +481,7 @@ def is_valid_zip(path):
 
 
 def find_existing_offline_zip():
-    search_dirs = [MODELSCOPE_DOWNLOAD_DIR, REPO_PATH]
+    search_dirs = [MODELSCOPE_DOWNLOAD_DIR, OFFLINE_BUILD_CACHE_DIR, REPO_PATH]
 
     for directory in search_dirs:
         if not os.path.isdir(directory):
@@ -518,75 +560,398 @@ def install_modelscope_cli():
     return False
 
 
-def download_offline_build_from_modelscope():
-    os.makedirs(MODELSCOPE_DOWNLOAD_DIR, exist_ok=True)
+def extract_urls_from_text(text):
+    if not text:
+        return []
+
+    urls = re.findall(r"https?://[^\s'\")<>]+", text)
+    cleaned_urls = []
+
+    for url in urls:
+        cleaned_urls.append(url.rstrip(".,;"))
+
+    return cleaned_urls
+
+
+def convert_internal_oss_url_to_public(url):
+    if not url:
+        return None
+
+    public_url = url.replace("-internal.aliyuncs.com", ".aliyuncs.com")
+    public_url = public_url.replace("http://", "https://", 1)
+
+    if public_url == url:
+        return None
+
+    return public_url
+
+
+def download_url_to_file(url, output_path, use_proxy=False):
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+
+    env = get_modelscope_download_env(use_proxy=use_proxy)
+
+    if shutil.which("aria2c"):
+        command = [
+            "aria2c",
+            "--continue=true",
+            "--max-connection-per-server=8",
+            "--split=8",
+            "--min-split-size=16M",
+            "--check-certificate=false",
+            "--summary-interval=30",
+            "--timeout=30",
+            "--retry-wait=10",
+            "--max-tries=20",
+            "--dir",
+            os.path.dirname(output_path),
+            "--out",
+            os.path.basename(output_path),
+            url
+        ]
+    else:
+        command = [
+            "wget",
+            "--continue",
+            "--timeout=30",
+            "--tries=20",
+            "--no-check-certificate",
+            "--output-document",
+            output_path,
+            url
+        ]
+
+    try:
+        print("执行直链下载命令: " + " ".join(command[:1]) + " ...")
+
+        result = subprocess.run(
+            command,
+            env=env,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            timeout=MODELSCOPE_DOWNLOAD_TIMEOUT
+        )
+
+        if result.stdout.strip():
+            print(result.stdout[-2000:])
+        if result.stderr.strip():
+            print(result.stderr[-2000:])
+
+        if result.returncode == 0 and is_valid_zip(output_path):
+            print(f"直链下载成功: {output_path}")
+            return True
+
+        print(f"直链下载失败或压缩包无效: {output_path}")
+        return False
+
+    except subprocess.TimeoutExpired:
+        print("直链下载超时。")
+        return False
+    except Exception as e:
+        print(f"直链下载异常: {e}")
+        return False
+
+
+def try_public_oss_fallback(download_output, zip_name):
+    urls = extract_urls_from_text(download_output)
+
+    for url in urls:
+        public_url = convert_internal_oss_url_to_public(url)
+        if not public_url:
+            continue
+
+        print("检测到 ModelScope 返回 OSS 内网地址，尝试改用公网 OSS 地址继续下载。")
+        target_path = os.path.join(MODELSCOPE_DOWNLOAD_DIR, zip_name)
+
+        # 先直连公网 OSS；失败后才走代理。直连更适合国内服务器访问阿里云 OSS。
+        if download_url_to_file(public_url, target_path, use_proxy=False):
+            return target_path
+
+        if download_url_to_file(public_url, target_path, use_proxy=True):
+            return target_path
+
+    return None
+
+
+def ensure_git_lfs_available():
+    try:
+        result = subprocess.run(
+            ["git", "lfs", "version"],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            timeout=20
+        )
+        if result.returncode == 0:
+            print("检测到 git-lfs。")
+            return True
+    except Exception:
+        pass
+
+    print("未检测到 git-lfs，尝试安装 git-lfs。")
+
+    install_commands = []
+    if shutil.which("apt-get"):
+        install_commands.append(["bash", "-lc", "apt-get update && apt-get install -y git-lfs"])
+    if shutil.which("yum"):
+        install_commands.append(["bash", "-lc", "yum install -y git-lfs"])
+
+    for command in install_commands:
+        try:
+            result = subprocess.run(
+                command,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+                timeout=600
+            )
+            if result.stdout.strip():
+                print(result.stdout[-2000:])
+            if result.stderr.strip():
+                print(result.stderr[-2000:])
+
+            check = subprocess.run(
+                ["git", "lfs", "version"],
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+                timeout=20
+            )
+            if check.returncode == 0:
+                print("git-lfs 安装成功。")
+                return True
+
+        except subprocess.TimeoutExpired:
+            print("git-lfs 安装超时。")
+        except Exception as e:
+            print(f"git-lfs 安装异常: {e}")
+
+    print("git-lfs 不可用，后续 git clone 可能只能拿到 LFS 指针文件。")
+    return False
+
+
+def clone_whole_forgeui_repo_with_git():
+    os.makedirs(OFFLINE_BUILD_CACHE_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():
+    if not shutil.which("git"):
+        print("未检测到 git，跳过 git 整库克隆。")
         return None
 
-    print(f"开始从 ModelScope 下载 {MODELSCOPE_MODEL_ID} 的 offline_build 压缩包。")
+    ensure_git_lfs_available()
+
+    repo_urls = [
+        "https://www.modelscope.cn/HOPStudio/forgeui.git",
+        "https://modelscope.cn/HOPStudio/forgeui.git"
+    ]
 
-    env = get_proxy_env()
-    download_commands = []
+    if os.path.isdir(os.path.join(MODELSCOPE_DOWNLOAD_DIR, ".git")):
+        print(f"检测到已存在 ModelScope 整库仓库: {MODELSCOPE_DOWNLOAD_DIR}")
+        try:
+            subprocess.run(
+                ["git", "-C", MODELSCOPE_DOWNLOAD_DIR, "lfs", "pull"],
+                env=get_modelscope_download_env(use_proxy=False),
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+                timeout=MODELSCOPE_DOWNLOAD_TIMEOUT
+            )
+        except Exception as e:
+            print(f"git lfs pull 执行异常: {e}")
 
-    # 优先按指定文件下载，避免把整个仓库全部拉下来。
-    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:
+        existing_zip = find_existing_offline_zip()
+        if existing_zip:
+            print(f"在已克隆仓库中找到 offline_build 压缩包: {existing_zip}")
+            return existing_zip
+
+    elif os.path.exists(MODELSCOPE_DOWNLOAD_DIR):
+        # 这里是旧版 modelscope download 留下的目录。整库 git clone 需要空目录，
+        # 因此先改名备份，不直接删除，避免误删已有断点文件。
+        backup_dir = MODELSCOPE_DOWNLOAD_DIR + f".bak_{int(time.time())}"
         try:
-            print("执行下载命令: " + " ".join(command))
+            shutil.move(MODELSCOPE_DOWNLOAD_DIR, backup_dir)
+            print(f"检测到非 git 下载目录，已备份为: {backup_dir}")
+        except Exception as e:
+            print(f"备份旧下载目录失败: {e}")
+            return None
+
+    for repo_url in repo_urls:
+        for use_proxy, mode_name in [(False, "直连"), (True, "代理")]:
+            if os.path.exists(MODELSCOPE_DOWNLOAD_DIR):
+                shutil.rmtree(MODELSCOPE_DOWNLOAD_DIR, ignore_errors=True)
+
+            env = get_modelscope_download_env(use_proxy=use_proxy)
+            command = [
+                "git",
+                "clone",
+                "--depth",
+                "1",
+                repo_url,
+                MODELSCOPE_DOWNLOAD_DIR
+            ]
+
+            try:
+                print(f"执行 ModelScope 整库 git clone [{mode_name}]: " + " ".join(command))
+
+                result = subprocess.run(
+                    command,
+                    env=env,
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.PIPE,
+                    text=True,
+                    timeout=MODELSCOPE_DOWNLOAD_TIMEOUT
+                )
+
+                if result.stdout.strip():
+                    print(result.stdout[-2000:])
+                if result.stderr.strip():
+                    print(result.stderr[-4000:])
+
+                if result.returncode != 0:
+                    print("本次 git clone 失败，继续尝试下一个方式。")
+                    continue
+
+                try:
+                    subprocess.run(
+                        ["git", "-C", MODELSCOPE_DOWNLOAD_DIR, "lfs", "pull"],
+                        env=env,
+                        stdout=subprocess.PIPE,
+                        stderr=subprocess.PIPE,
+                        text=True,
+                        timeout=MODELSCOPE_DOWNLOAD_TIMEOUT
+                    )
+                except Exception as e:
+                    print(f"git lfs pull 执行异常: {e}")
+
+                zip_path = find_existing_offline_zip()
+                if zip_path:
+                    print(f"ModelScope 整库克隆完成，找到 offline_build 压缩包: {zip_path}")
+                    return zip_path
+
+                print("整库 git clone 完成，但未找到有效 offline_build 压缩包。")
+
+            except subprocess.TimeoutExpired:
+                print("ModelScope 整库 git clone 超时，继续尝试下一个方式。")
+            except Exception as e:
+                print(f"ModelScope 整库 git clone 异常: {e}")
+
+    return None
+
+
+def download_whole_forgeui_repo_with_modelscope_cli():
+    if not install_modelscope_cli():
+        return None
+
+    commands = [
+        (
+            [
+                "modelscope",
+                "download",
+                "--model",
+                MODELSCOPE_MODEL_ID,
+                "--local_dir",
+                MODELSCOPE_DOWNLOAD_DIR
+            ],
+            False,
+            "整库直连"
+        ),
+        (
+            [
+                "modelscope",
+                "download",
+                "--model",
+                MODELSCOPE_MODEL_ID,
+                "--local_dir",
+                MODELSCOPE_DOWNLOAD_DIR,
+                "--endpoint",
+                "https://www.modelscope.cn"
+            ],
+            False,
+            "整库直连+显式 endpoint"
+        ),
+        (
+            [
+                "modelscope",
+                "download",
+                "--model",
+                MODELSCOPE_MODEL_ID,
+                "--local_dir",
+                MODELSCOPE_DOWNLOAD_DIR
+            ],
+            True,
+            "整库代理"
+        )
+    ]
+
+    for command, use_proxy, mode_name in commands:
+        try:
+            print(f"执行 ModelScope CLI 整库下载 [{mode_name}]: " + " ".join(command))
 
             result = subprocess.run(
                 command,
-                env=env,
+                env=get_modelscope_download_env(use_proxy=use_proxy),
                 stdout=subprocess.PIPE,
                 stderr=subprocess.PIPE,
                 text=True,
-                timeout=3600
+                timeout=MODELSCOPE_DOWNLOAD_TIMEOUT
             )
 
+            output_text = ""
             if result.stdout.strip():
+                output_text += result.stdout + "\n"
                 print(result.stdout[-2000:])
             if result.stderr.strip():
-                print(result.stderr[-2000:])
+                output_text += result.stderr + "\n"
+                print(result.stderr[-4000:])
 
             zip_path = find_existing_offline_zip()
             if result.returncode == 0 and zip_path:
-                print(f"offline_build 压缩包下载成功: {zip_path}")
+                print(f"ModelScope CLI 整库下载完成，找到 offline_build 压缩包: {zip_path}")
                 return zip_path
 
-            print("本次下载未找到有效 offline_build 压缩包，继续尝试下一个命令。")
+            # 这里不再按文件名调用 modelscope download，只在 CLI 日志暴露 OSS 直链时做公网地址回退。
+            fallback_zip = try_public_oss_fallback(output_text, "offlinbuild.zip")
+            if fallback_zip and is_valid_zip(fallback_zip):
+                print(f"offline_build 压缩包通过公网 OSS 回退下载成功: {fallback_zip}")
+                return fallback_zip
+
+            print("本次整库下载未找到有效 offline_build 压缩包，继续尝试下一个命令。")
 
         except subprocess.TimeoutExpired:
-            print("ModelScope 下载超时，继续尝试下一个命令。")
+            print("ModelScope CLI 整库下载超时，继续尝试下一个命令。")
         except Exception as e:
-            print(f"ModelScope 下载异常: {e}")
+            print(f"ModelScope CLI 整库下载异常: {e}")
+
+    return None
+
+
+def download_offline_build_from_modelscope():
+    os.makedirs(MODELSCOPE_DOWNLOAD_DIR, exist_ok=True)
+    os.makedirs(MODELSCOPE_CACHE_DIR, exist_ok=True)
+
+    existing_zip = find_existing_offline_zip()
+    if existing_zip:
+        print(f"检测到已有 offline_build 压缩包: {existing_zip}")
+        return existing_zip
+
+    print(f"开始从 ModelScope 整库克隆/下载 {MODELSCOPE_MODEL_ID}。")
+
+    zip_path = clone_whole_forgeui_repo_with_git()
+    if zip_path:
+        return zip_path
+
+    print("git 整库克隆未成功，改用 ModelScope CLI 整库下载。")
+
+    zip_path = download_whole_forgeui_repo_with_modelscope_cli()
+    if zip_path:
+        return zip_path
 
-    print("ModelScope 下载失败，未找到有效 offline_build 压缩包。")
+    print("ModelScope 整库克隆/下载失败，未找到有效 offline_build 压缩包。")
     return None