main.py

import os
import subprocess
import threading
import shutil
import json
import requests
import gradio as gr


REPO_URL = "https://hf-mirror.com/HOPStudio/modelscope"
REPO_PATH = "/root/modelscope"

FRP_VERSION = "0.69.0"
CURRENT_DIR = REPO_PATH
LAST_OUTPUT = ""


def clone_repository():
    if os.path.exists(REPO_PATH):
        print(f"检测到已有仓库，正在删除: {REPO_PATH}")
        try:
            shutil.rmtree(REPO_PATH)
            print("旧仓库已删除。")
        except Exception as e:
            print(f"删除旧仓库失败: {e}")
            return False

    try:
        print(f"正在克隆仓库到 {REPO_PATH}...")

        env = os.environ.copy()
        env["GIT_LFS_SKIP_SMUDGE"] = "1"
        env["HF_HUB_DISABLE_XET"] = "1"

        subprocess.run(
            ["git", "clone", REPO_URL, REPO_PATH],
            check=True,
            env=env,
            timeout=120
        )

        print("仓库克隆成功。")
        return True

    except subprocess.TimeoutExpired:
        print("克隆仓库超时。")
        return False

    except subprocess.CalledProcessError as e:
        print(f"克隆仓库失败: {e}")
        return False

    except Exception as e:
        print(f"克隆仓库异常: {e}")
        return False


def execute_shell_command(command):
    global CURRENT_DIR

    command = command.strip()

    if command == "pwd":
        return CURRENT_DIR

    if command.startswith("cd "):
        path = command[3:].strip()
        path = os.path.expanduser(path)

        if not os.path.isabs(path):
            path = os.path.join(CURRENT_DIR, path)

        path = os.path.abspath(path)

        if os.path.isdir(path):
            CURRENT_DIR = path
            return f"已切换目录: {CURRENT_DIR}"

        return f"目录不存在: {path}"

    try:
        result = subprocess.run(
            command,
            shell=True,
            cwd=CURRENT_DIR,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            timeout=120
        )

        output = result.stdout + result.stderr

        if not output.strip():
            output = "命令执行完成，无输出。"

        return f"当前目录: {CURRENT_DIR}\n\n{output}"

    except subprocess.TimeoutExpired:
        return "命令执行超时。"

    except Exception as e:
        return f"执行命令失败: {e}"


def _api_call_impl(input_text, model="internlm2.5-latest", temperature=0.8, top_p=0.9, n=1):
    if input_text.startswith("cmd_run "):
        command = input_text[len("cmd_run "):].strip()
        return execute_shell_command(command)

    api_token = os.getenv("API_TOKEN")

    if not api_token:
        return "API_TOKEN 未设置。"

    url = "https://internlm-chat.intern-ai.org.cn/puyu/api/v1/chat/completions"

    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_token}"
    }

    data = {
        "model": model,
        "messages": [
            {
                "role": "user",
                "content": input_text
            }
        ],
        "n": int(n),
        "temperature": float(temperature),
        "top_p": float(top_p)
    }

    try:
        response = requests.post(
            url,
            headers=headers,
            data=json.dumps(data),
            timeout=60
        )

        response.raise_for_status()
        return response.json()["choices"][0]["message"]["content"]

    except Exception as e:
        return f"API 请求失败: {e}"


def api_call(input_text, model="internlm2.5-latest", temperature=0.8, top_p=0.9, n=1):
    global LAST_OUTPUT

    LAST_OUTPUT = _api_call_impl(
        input_text,
        model,
        temperature,
        top_p,
        n
    )

    return LAST_OUTPUT


def refresh_frontend():
    return LAST_OUTPUT


def create_frpc_toml():
    os.makedirs(REPO_PATH, exist_ok=True)

    frpc_toml_path = os.path.join(REPO_PATH, "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"
localIP = "127.0.0.1"
localPort = 8188
remotePort = 8188
"""

    try:
        with open(frpc_toml_path, "w", encoding="utf-8") as f:
            f.write(toml_content)

        print(f"frpc.toml 已创建: {frpc_toml_path}")

    except Exception as e:
        print(f"创建 frpc.toml 失败: {e}")


def install_frp():
    frp_tar_name = f"frp_{FRP_VERSION}_linux_amd64.tar.gz"
    frp_dir_name = f"frp_{FRP_VERSION}_linux_amd64"

    frp_tar_path = os.path.join(REPO_PATH, frp_tar_name)
    frp_install_path = os.path.join(REPO_PATH, frp_dir_name)

    frp_download_urls = [
        f"https://github.com/fatedier/frp/releases/download/v{FRP_VERSION}/{frp_tar_name}",
        f"https://gh.llkk.cc/https://github.com/fatedier/frp/releases/download/v{FRP_VERSION}/{frp_tar_name}",
        f"https://gh-proxy.com/https://github.com/fatedier/frp/releases/download/v{FRP_VERSION}/{frp_tar_name}"
    ]

    if os.path.exists(frp_install_path):
        print(f"frp 已安装: {frp_install_path}")
        return True

    def is_valid_tar_gz(path):
        if not os.path.exists(path):
            return False

        try:
            result = subprocess.run(
                ["tar", "-tzf", path],
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                timeout=20
            )
            return result.returncode == 0
        except Exception:
            return False

    if os.path.exists(frp_tar_path) and not is_valid_tar_gz(frp_tar_path):
        print("检测到损坏的 frp 压缩包，正在删除。")
        os.remove(frp_tar_path)

    if not os.path.exists(frp_tar_path):
        print("开始下载 frp。")

        download_success = False

        for url in frp_download_urls:
            try:
                print(f"尝试下载: {url}")

                result = subprocess.run(
                    [
                        "wget",
                        "--timeout=30",
                        "--tries=1",
                        "--no-check-certificate",
                        "-O",
                        frp_tar_path,
                        url
                    ],
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                    text=True,
                    timeout=90
                )

                if result.returncode == 0 and is_valid_tar_gz(frp_tar_path):
                    print("frp 下载成功。")
                    download_success = True
                    break

                print("下载失败或文件无效。")

                if os.path.exists(frp_tar_path):
                    os.remove(frp_tar_path)

            except subprocess.TimeoutExpired:
                print("frp 下载超时。")

                if os.path.exists(frp_tar_path):
                    os.remove(frp_tar_path)

            except Exception as e:
                print(f"frp 下载异常: {e}")

                if os.path.exists(frp_tar_path):
                    os.remove(frp_tar_path)

        if not download_success:
            print("frp 下载失败，跳过 frp。")
            return False

    try:
        print(f"正在解压 frp 到 {REPO_PATH}...")

        subprocess.run(
            ["tar", "-xzf", frp_tar_path, "-C", REPO_PATH],
            check=True,
            timeout=30
        )

        frpc_path = os.path.join(frp_install_path, "frpc")
        frps_path = os.path.join(frp_install_path, "frps")

        if os.path.exists(frpc_path):
            subprocess.run(["chmod", "+x", frpc_path], check=False)

        if os.path.exists(frps_path):
            subprocess.run(["chmod", "+x", frps_path], check=False)

        print("frp 安装成功。")
        return True

    except subprocess.TimeoutExpired:
        print("frp 解压超时。")
        return False

    except subprocess.CalledProcessError as e:
        print(f"frp 解压失败: {e}")
        return False

    except Exception as e:
        print(f"frp 安装异常: {e}")
        return False


def start_frp():
    frp_dir_name = f"frp_{FRP_VERSION}_linux_amd64"

    frpc_path = os.path.join(REPO_PATH, frp_dir_name, "frpc")
    frpc_config_path = os.path.join(REPO_PATH, "frpc.toml")

    if not os.path.exists(frpc_path):
        print(f"frpc 不存在: {frpc_path}")
        return

    if not os.path.exists(frpc_config_path):
        create_frpc_toml()

    try:
        subprocess.Popen(
            [frpc_path, "-c", frpc_config_path],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL
        )

        print("frp 已启动。")

    except Exception as e:
        print(f"启动 frp 失败: {e}")


def start_jupyterlab_process():
    jupyter_script = os.path.join(REPO_PATH, "start_jupyterlab.py")

    if not os.path.exists(jupyter_script):
        print(f"未找到 JupyterLab 启动脚本: {jupyter_script}")
        return

    try:
        subprocess.Popen(
            ["python3", jupyter_script],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL
        )

        print("JupyterLab 已启动。")

    except Exception as e:
        print(f"启动 JupyterLab 失败: {e}")


# ===== 新增：ModelScope 下载 comfy_portable.zip，解压后使用自带 Python 启动 ComfyUI =====
MODELSCOPE_MODEL_ID = "HOPStudio/forgeui"
COMFY_PORTABLE_ZIP_NAME = "comfy_portable.zip"
COMFY_PORTABLE_ZIP_PATH = os.path.join("/root", COMFY_PORTABLE_ZIP_NAME)
COMFY_PORTABLE_DIR = "/root/comfy_portable"
COMFY_PORTABLE_PYTHON = os.path.join(COMFY_PORTABLE_DIR, "python", "bin", "python3")
COMFYUI_MAIN = os.path.join(COMFY_PORTABLE_DIR, "ComfyUI", "main.py")
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


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 get_modelscope_env(use_proxy=False):
    env = os.environ.copy()
    env["MODELSCOPE_ENDPOINT"] = "https://www.modelscope.cn"
    env["MODELSCOPE_SDK_ENDPOINT"] = "https://www.modelscope.cn"

    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)

    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:
        import zipfile
        if not zipfile.is_zipfile(path):
            return False

        with zipfile.ZipFile(path, "r") as zf:
            return len(zf.namelist()) > 0
    except Exception:
        return False


def is_comfy_portable_ready():
    return os.path.exists(COMFY_PORTABLE_PYTHON) and os.path.exists(COMFYUI_MAIN)


def ensure_comfy_portable_permissions():
    """修复 zip 解压后丢失的可执行权限，只作用于 comfy_portable。"""
    if not os.path.isdir(COMFY_PORTABLE_DIR):
        return

    paths_to_chmod = [
        COMFY_PORTABLE_PYTHON,
        os.path.join(COMFY_PORTABLE_DIR, "python", "bin", "python"),
        os.path.join(COMFY_PORTABLE_DIR, "python", "bin", "python3.10"),
        os.path.join(COMFY_PORTABLE_DIR, "start.sh"),
    ]

    for path in paths_to_chmod:
        if not os.path.exists(path):
            continue

        try:
            current_mode = os.stat(path).st_mode
            os.chmod(path, current_mode | 0o755)
            print(f"已修复执行权限: {path}")
        except Exception as e:
            print(f"修复执行权限失败: {path}, {e}")

    python_bin_dir = os.path.join(COMFY_PORTABLE_DIR, "python", "bin")
    if os.path.isdir(python_bin_dir):
        try:
            subprocess.run(
                [
                    "bash",
                    "-lc",
                    f"find {python_bin_dir!r} -type f -exec chmod +x {{}} \\;"
                ],
                check=False,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
                timeout=60
            )
            print(f"已修复 python/bin 下文件执行权限: {python_bin_dir}")
        except Exception as e:
            print(f"批量修复 python/bin 权限失败: {e}")


def download_comfy_portable_from_modelscope():
    if is_comfy_portable_ready():
        print(f"检测到已解压的 comfy_portable: {COMFY_PORTABLE_DIR}")
        return True

    if is_valid_zip(COMFY_PORTABLE_ZIP_PATH):
        print(f"检测到已有 comfy_portable 压缩包: {COMFY_PORTABLE_ZIP_PATH}")
        return True

    if os.path.exists(COMFY_PORTABLE_ZIP_PATH):
        print(f"检测到无效 comfy_portable 压缩包，正在删除: {COMFY_PORTABLE_ZIP_PATH}")
        try:
            os.remove(COMFY_PORTABLE_ZIP_PATH)
        except Exception as e:
            print(f"删除无效压缩包失败: {e}")
            return False

    if not shutil.which("modelscope"):
        print("未检测到 modelscope 命令。为避免改动系统 Python 依赖，本脚本不会自动安装 modelscope。")
        print("请先在外部环境安装好 modelscope CLI，或手动把 comfy_portable.zip 放到 /root/comfy_portable.zip。")
        return False

    commands = [
        (False, "直连"),
        (True, "代理")
    ]

    for use_proxy, mode_name in commands:
        command = [
            "modelscope",
            "download",
            "--model",
            MODELSCOPE_MODEL_ID,
            COMFY_PORTABLE_ZIP_NAME,
            "--local_dir",
            "/root"
        ]

        try:
            print(f"执行 ModelScope 下载 comfy_portable.zip [{mode_name}]: " + " ".join(command))

            result = subprocess.run(
                command,
                env=get_modelscope_env(use_proxy=use_proxy),
                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 and is_valid_zip(COMFY_PORTABLE_ZIP_PATH):
                print(f"comfy_portable.zip 下载完成: {COMFY_PORTABLE_ZIP_PATH}")
                return True

            print("本次 ModelScope 下载失败或压缩包无效，继续尝试下一个方式。")

        except subprocess.TimeoutExpired:
            print("ModelScope 下载 comfy_portable.zip 超时，继续尝试下一个方式。")
        except Exception as e:
            print(f"ModelScope 下载 comfy_portable.zip 异常: {e}")

    print("ModelScope 下载 comfy_portable.zip 失败。")
    return False


def extract_comfy_portable():
    if is_comfy_portable_ready():
        print(f"comfy_portable 已存在，跳过解压: {COMFY_PORTABLE_DIR}")
        return True

    if not is_valid_zip(COMFY_PORTABLE_ZIP_PATH):
        print(f"comfy_portable.zip 不存在或无效: {COMFY_PORTABLE_ZIP_PATH}")
        return False

    if os.path.exists(COMFY_PORTABLE_DIR):
        print(f"检测到不完整的 comfy_portable 目录，正在删除后重新解压: {COMFY_PORTABLE_DIR}")
        try:
            shutil.rmtree(COMFY_PORTABLE_DIR)
        except Exception as e:
            print(f"删除不完整 comfy_portable 目录失败: {e}")
            return False

    try:
        print(f"正在解压 comfy_portable.zip 到 /root: {COMFY_PORTABLE_ZIP_PATH}")

        import zipfile
        with zipfile.ZipFile(COMFY_PORTABLE_ZIP_PATH, "r") as zf:
            zf.extractall("/root")

        ensure_comfy_portable_permissions()

        if is_comfy_portable_ready():
            print(f"comfy_portable 解压完成: {COMFY_PORTABLE_DIR}")
            return True

        print("解压完成，但未找到 /root/comfy_portable/python/bin/python3 或 /root/comfy_portable/ComfyUI/main.py。")
        return False

    except Exception as e:
        print(f"解压 comfy_portable.zip 失败: {e}")
        return False


def prepare_comfy_portable():
    if is_comfy_portable_ready():
        print(f"comfy_portable 已准备完成: {COMFY_PORTABLE_DIR}")
        return True

    if not download_comfy_portable_from_modelscope():
        return False

    return extract_comfy_portable()


def start_comfyui_process():
    if not prepare_comfy_portable():
        print("comfy_portable 准备失败，跳过 ComfyUI 启动。")
        return

    ensure_comfy_portable_permissions()

    if not os.path.exists(COMFY_PORTABLE_PYTHON):
        print(f"未找到便携 Python: {COMFY_PORTABLE_PYTHON}")
        return

    if not os.path.exists(COMFYUI_MAIN):
        print(f"未找到 ComfyUI main.py: {COMFYUI_MAIN}")
        return

    configure_git_proxy()

    python_lib_dir = os.path.join(COMFY_PORTABLE_DIR, "python", "lib")
    comfyui_dir = os.path.join(COMFY_PORTABLE_DIR, "ComfyUI")

    command = f"""
set -e
export LD_LIBRARY_PATH=\"{python_lib_dir}:${{LD_LIBRARY_PATH:-}}\"
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 \"{comfyui_dir}\"
\"{COMFY_PORTABLE_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=COMFY_PORTABLE_DIR,
            env=get_proxy_env(),
            stdout=log_file,
            stderr=subprocess.STDOUT
        )

        import time
        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 控制台")

        input_box = gr.Textbox(
            label="输入",
            placeholder="输入对话内容"
        )

        model_box = gr.Textbox(
            value="internlm2.5-latest",
            label="模型"
        )

        temperature_slider = gr.Slider(
            minimum=0,
            maximum=1,
            value=0.8,
            step=0.01,
            label="Temperature"
        )

        top_p_slider = gr.Slider(
            minimum=0,
            maximum=1,
            value=0.9,
            step=0.01,
            label="Top-p"
        )

        n_slider = gr.Slider(
            minimum=1,
            maximum=10,
            value=1,
            step=1,
            label="生成回复数量"
        )

        with gr.Row():
            send_button = gr.Button("发送", variant="primary")
            refresh_button = gr.Button("刷新前端")

        output_box = gr.Textbox(
            label="回复",
            lines=20,
            interactive=False
        )

        send_button.click(
            api_call,
            inputs=[
                input_box,
                model_box,
                temperature_slider,
                top_p_slider,
                n_slider
            ],
            outputs=output_box
        )

        refresh_button.click(
            refresh_frontend,
            outputs=output_box
        )

        timer = gr.Timer(5)

        timer.tick(
            refresh_frontend,
            outputs=output_box
        )

    print("Gradio 正在启动。")

    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )


def main():
    global CURRENT_DIR

    clone_success = clone_repository()

    if clone_success:
        CURRENT_DIR = REPO_PATH

    install_frp()
    create_frpc_toml()

    try:
        frp_thread = threading.Thread(
            target=start_frp,
            daemon=True
        )
        frp_thread.start()
    except Exception as e:
        print(f"frp 线程启动失败: {e}")

    try:
        jupyter_thread = threading.Thread(
            target=start_jupyterlab_process,
            daemon=True
        )
        jupyter_thread.start()
    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()


if __name__ == "__main__":
    main()