69 lines
2.2 KiB
Python
69 lines
2.2 KiB
Python
"""PP-OCRv6 unified launcher using isolated CPU/GPU uv projects."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import shutil
|
||
import subprocess
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
ROOT = Path(__file__).resolve().parent
|
||
|
||
|
||
def _requested_device(argv: list[str]) -> str:
|
||
for index, value in enumerate(argv):
|
||
if value == "--device":
|
||
if index + 1 >= len(argv):
|
||
raise SystemExit("--device 需要 cpu 或 gpu")
|
||
return argv[index + 1].lower()
|
||
if value.startswith("--device="):
|
||
return value.split("=", 1)[1].lower()
|
||
return "cpu"
|
||
|
||
|
||
def main() -> int:
|
||
try:
|
||
device = _requested_device(sys.argv[1:])
|
||
except SystemExit as exc:
|
||
print(exc, file=sys.stderr)
|
||
return 2
|
||
if device not in {"cpu", "gpu"}:
|
||
print(f"不支持的设备: {device}，可选 cpu/gpu", file=sys.stderr)
|
||
return 2
|
||
|
||
project = ROOT / device
|
||
runner = project / "runner.py"
|
||
if sys.platform == "win32":
|
||
environment_python = project / ".venv" / "Scripts" / "python.exe"
|
||
else:
|
||
environment_python = project / ".venv" / "bin" / "python"
|
||
|
||
if device == "gpu":
|
||
ready_marker = project / ".gpu-ready"
|
||
if not ready_marker.is_file() or not environment_python.is_file():
|
||
print(
|
||
"GPU 环境尚未安装完成。请根据目标 CUDA 版本运行：\n"
|
||
" python gpu/setup_env.py --cuda cu118\n"
|
||
"或：\n"
|
||
" python gpu/setup_env.py --cuda cu126",
|
||
file=sys.stderr,
|
||
)
|
||
return 1
|
||
command = [str(environment_python), str(runner), *sys.argv[1:]]
|
||
return subprocess.run(command, cwd=ROOT).returncode
|
||
|
||
if environment_python.is_file():
|
||
command = [str(environment_python), str(runner), *sys.argv[1:]]
|
||
return subprocess.run(command, cwd=ROOT).returncode
|
||
|
||
uv = shutil.which("uv")
|
||
if not uv:
|
||
print("未找到 CPU 虚拟环境和 uv，请先安装 uv。", file=sys.stderr)
|
||
return 1
|
||
command = [uv, "run", "--project", str(project), "python", str(runner), *sys.argv[1:]]
|
||
return subprocess.run(command, cwd=ROOT).returncode
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|