scripts/launch_minicpm_omni.py

from __future__ import annotations

import json
import os
import shutil
import signal
import subprocess
import sys
import time
import urllib.parse
from pathlib import Path


def main() -> int:
    root = Path(os.environ.get("MINICPM_COMNI_ROOT", "")).expanduser().resolve()
    if not root.exists() or not (root / "worker.py").exists() or not (root / "gateway.py").exists():
        print("MINICPM_COMNI_ROOT must point to an OpenBMB/MiniCPM-o-Demo Comni checkout.", file=sys.stderr)
        return 2

    model_dir = Path(os.environ.get("MINICPM_MODEL_DIR", "")).expanduser().resolve()
    llm_model = os.environ.get("MINICPM_LLM_MODEL", "").strip()
    llama_root = Path(os.environ.get("MINICPM_LLAMACPP_ROOT", "")).expanduser().resolve()
    context = int(os.environ.get("MINICPM_CTX_SIZE", "8192"))
    gpu_layers_raw = os.environ.get("MINICPM_N_GPU_LAYERS", "auto").strip().lower()
    gpu_layers = 99 if gpu_layers_raw == "auto" else int(gpu_layers_raw)
    gateway_url = os.environ.get("MINICPM_GATEWAY_URL", "http://127.0.0.1:8006")
    gateway_port = urllib.parse.urlparse(gateway_url).port or 8006
    worker_port = gateway_port + 1
    _stop_stale_internal_server(19060, llama_root)

    if not (model_dir / llm_model).exists():
        print(f"Selected model does not exist: {model_dir / llm_model}", file=sys.stderr)
        return 2
    if not llama_root.exists():
        print(f"llama.cpp-omni root does not exist: {llama_root}", file=sys.stderr)
        return 2

    config_path = root / "config.json"
    example_config_path = root / "config.example.json"
    seed_config_path = config_path if config_path.exists() else example_config_path
    config = json.loads(seed_config_path.read_text(encoding="utf-8")) if seed_config_path.exists() else {}
    config["backend"] = "cpp"
    config.setdefault("model", {})["model_path"] = str(model_dir)
    config.setdefault("service", {}).update({
        "gateway_port": gateway_port,
        "worker_base_port": worker_port,
        "num_workers": 1,
    })
    config.setdefault("cpp_backend", {}).update({
        "llamacpp_root": str(llama_root),
        "model_dir": str(model_dir),
        "llm_model": llm_model,
        "ctx_size": context,
        "n_gpu_layers": gpu_layers,
    })
    config_path.write_text(json.dumps(config, indent=2) + "\n", encoding="utf-8")

    python = _external_python(root)
    log_dir = root / "tmp"
    log_dir.mkdir(exist_ok=True)
    env = os.environ.copy()
    env["PYTHONPATH"] = str(root)
    gpu_device = os.environ.get("MINICPM_GPU_DEVICE", "auto").strip().lower()

    # MiniCPM-o-Demo's cpp_backend defaults to tts_gpu_layers=100 and
    # token2wav_device=gpu:0. On a 6 GB card the main LLM already fills VRAM,
    # so loading TTS on GPU OOMs and crashes llama-server during omni_init.
    # If the user hasn't picked an explicit policy, probe the GPU and keep
    # TTS+T2W on CPU whenever the device has less than ~8 GB free — the
    # text-generation hot path still runs on GPU, which is what matters.
    if "MINICPM_TTS_GPU_LAYERS" not in env or "MINICPM_TOKEN2WAV_DEVICE" not in env:
        keep_tts_on_cpu = _tts_should_use_cpu(gpu_device, gpu_layers)
        env.setdefault("MINICPM_TTS_GPU_LAYERS", "0" if keep_tts_on_cpu else "100")
        env.setdefault("MINICPM_TOKEN2WAV_DEVICE", "cpu" if keep_tts_on_cpu else "gpu:0")
    # Honor the picker's device choice. CPU and zero-layer modes blank every
    # visibility var so llama.cpp falls back to host execution regardless of
    # which backend was compiled in. cuda:N / rocm:N pin the chosen index.
    if gpu_device == "cpu" or gpu_layers == 0:
        env["CUDA_VISIBLE_DEVICES"] = ""
        env["HIP_VISIBLE_DEVICES"] = ""
        env["ROCR_VISIBLE_DEVICES"] = ""
    elif gpu_device.startswith("cuda:"):
        env["CUDA_VISIBLE_DEVICES"] = gpu_device.split(":", 1)[1]
    elif gpu_device.startswith("rocm:"):
        index = gpu_device.split(":", 1)[1]
        env["HIP_VISIBLE_DEVICES"] = index
        env["ROCR_VISIBLE_DEVICES"] = index

    worker_log = (log_dir / "phantom_grid_worker.log").open("a", encoding="utf-8")
    gateway_log = (log_dir / "phantom_grid_gateway.log").open("a", encoding="utf-8")
    worker = subprocess.Popen(
        [str(python), "worker.py", "--port", str(worker_port), "--gpu-id", "0", "--worker-index", "0"],
        cwd=root, env=env, stdout=worker_log, stderr=subprocess.STDOUT,
        creationflags=subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0,
    )
    gateway = subprocess.Popen(
        [str(python), "gateway.py", "--port", str(gateway_port), "--workers", f"127.0.0.1:{worker_port}", "--http"],
        cwd=root, env=env, stdout=gateway_log, stderr=subprocess.STDOUT,
        creationflags=subprocess.CREATE_NO_WINDOW if os.name == "nt" else 0,
    )
    children = [gateway, worker]

    def stop_children(*_: object) -> None:
        for child in children:
            if child.poll() is None:
                child.terminate()
        deadline = time.time() + 8
        for child in children:
            try:
                child.wait(timeout=max(0.1, deadline - time.time()))
            except subprocess.TimeoutExpired:
                child.kill()
        raise SystemExit(0)

    signal.signal(signal.SIGINT, stop_children)
    signal.signal(signal.SIGTERM, stop_children)
    print(f"MiniCPM-o worker PID {worker.pid}; gateway PID {gateway.pid}", flush=True)
    while all(child.poll() is None for child in children):
        time.sleep(1)
    stop_children()
    return 1


def _tts_should_use_cpu(gpu_device: str, gpu_layers: int) -> bool:
    # CPU is the right place for TTS when there's no GPU offload at all.
    if gpu_device == "cpu" or gpu_layers == 0:
        return True
    # Probe NVIDIA total VRAM. <8 GB → main LLM fills it; TTS must run on CPU.
    nvidia_smi = shutil.which("nvidia-smi")
    if not nvidia_smi:
        return False
    try:
        completed = subprocess.run(
            [nvidia_smi, "--query-gpu=memory.total", "--format=csv,noheader,nounits"],
            capture_output=True, text=True, timeout=4, check=False,
        )
    except (OSError, subprocess.TimeoutExpired):
        return False
    if completed.returncode != 0:
        return False
    try:
        total_mib = max(int(line.strip()) for line in completed.stdout.splitlines() if line.strip())
    except ValueError:
        return False
    return total_mib < 8192


def _external_python(root: Path) -> Path:
    configured = os.environ.get("MINICPM_COMNI_PYTHON")
    candidates = [
        Path(configured) if configured else None,
        root / ".venv" / "base" / "Scripts" / "python.exe",
        root / ".venv" / "base" / "bin" / "python",
        root / ".venv" / "Scripts" / "python.exe",
        root / ".venv" / "bin" / "python",
    ]
    for candidate in candidates:
        if candidate and candidate.exists():
            return candidate
    return Path(sys.executable)


def _stop_stale_internal_server(port: int, llama_root: Path) -> None:
    """Clear an orphaned llama-server left by an interrupted Comni worker."""
    if os.name != "nt":
        return
    escaped_root = str(llama_root).replace("'", "''")
    script = (
        f"$root=[IO.Path]::GetFullPath('{escaped_root}'); "
        f"Get-NetTCPConnection -State Listen -LocalPort {port} -ErrorAction SilentlyContinue | "
        "ForEach-Object { $p=Get-CimInstance Win32_Process -Filter \"ProcessId=$($_.OwningProcess)\"; "
        "if ($p -and $p.Name -eq 'llama-server.exe' -and $p.ExecutablePath -and "
        "[IO.Path]::GetFullPath($p.ExecutablePath).StartsWith($root, [StringComparison]::OrdinalIgnoreCase)) "
        "{ Stop-Process -Id $p.ProcessId -Force } }"
    )
    try:
        subprocess.run(
            ["powershell", "-NoProfile", "-Command", script],
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
            timeout=10,
            check=False,
        )
    except (OSError, subprocess.TimeoutExpired):
        pass


if __name__ == "__main__":
    raise SystemExit(main())