wehub-resource-sync
173 lines
5.1 KiB
Python
173 lines
5.1 KiB
Python
import argparse
|
|
import ctypes
|
|
import glob
|
|
import os
|
|
import sys
|
|
import threading
|
|
from os.path import join
|
|
from dotenv import load_dotenv
|
|
|
|
DEFAULT_LEON_PROFILE = "just-me"
|
|
DEFAULT_TCP_SERVER_HOST = "127.0.0.1"
|
|
DEFAULT_TCP_SERVER_PORT = 5_367
|
|
|
|
|
|
def resolve_leon_home() -> str:
|
|
configured_leon_home = os.getenv("LEON_HOME", "").strip()
|
|
|
|
if configured_leon_home:
|
|
return os.path.abspath(configured_leon_home)
|
|
|
|
return os.path.join(os.path.expanduser("~"), ".leon")
|
|
|
|
|
|
def resolve_leon_profile() -> str:
|
|
return os.getenv("LEON_PROFILE", "").strip() or DEFAULT_LEON_PROFILE
|
|
|
|
|
|
def resolve_leon_profile_path() -> str:
|
|
configured_profile_path = os.getenv("LEON_PROFILE_PATH", "").strip()
|
|
|
|
if configured_profile_path:
|
|
return os.path.abspath(configured_profile_path)
|
|
|
|
return os.path.join(resolve_leon_home(), "profiles", resolve_leon_profile())
|
|
|
|
|
|
def _resolve_torch_root(pytorch_path: str) -> str | None:
|
|
normalized_path = os.path.abspath(pytorch_path)
|
|
if os.path.basename(normalized_path) == "torch" and os.path.isfile(
|
|
os.path.join(normalized_path, "__init__.py")
|
|
):
|
|
return normalized_path
|
|
|
|
torch_candidate = os.path.join(normalized_path, "torch")
|
|
if os.path.isfile(os.path.join(torch_candidate, "__init__.py")):
|
|
return torch_candidate
|
|
|
|
torch_nested_candidate = os.path.join(normalized_path, "torch", "torch")
|
|
if os.path.isfile(os.path.join(torch_nested_candidate, "__init__.py")):
|
|
return torch_nested_candidate
|
|
|
|
return None
|
|
|
|
|
|
def _add_pytorch_path(pytorch_path: str | None) -> str | None:
|
|
if not pytorch_path:
|
|
return None
|
|
|
|
torch_root = _resolve_torch_root(pytorch_path)
|
|
if torch_root:
|
|
sys.path.insert(0, os.path.dirname(torch_root))
|
|
return torch_root
|
|
|
|
sys.path.insert(0, os.path.abspath(pytorch_path))
|
|
return None
|
|
|
|
|
|
def _set_library_paths(paths: list[str]) -> None:
|
|
if not paths:
|
|
return
|
|
|
|
existing_path = ""
|
|
|
|
if sys.platform.startswith("win"):
|
|
add_dll_directory = getattr(os, "add_dll_directory", None)
|
|
for path in paths:
|
|
if os.path.isdir(path) and add_dll_directory:
|
|
add_dll_directory(path)
|
|
existing_path = os.environ.get("PATH", "")
|
|
os.environ["PATH"] = (
|
|
os.pathsep.join([*paths, existing_path])
|
|
if existing_path
|
|
else os.pathsep.join(paths)
|
|
)
|
|
return
|
|
|
|
if sys.platform == "darwin":
|
|
existing_path = os.environ.get("DYLD_LIBRARY_PATH", "")
|
|
os.environ["DYLD_LIBRARY_PATH"] = (
|
|
os.pathsep.join([*paths, existing_path])
|
|
if existing_path
|
|
else os.pathsep.join(paths)
|
|
)
|
|
return
|
|
|
|
existing_path = os.environ.get("LD_LIBRARY_PATH", "")
|
|
os.environ["LD_LIBRARY_PATH"] = (
|
|
os.pathsep.join([*paths, existing_path])
|
|
if existing_path
|
|
else os.pathsep.join(paths)
|
|
)
|
|
|
|
|
|
def _configure_external_libraries(
|
|
pytorch_path: str | None, nvidia_path: str | None
|
|
) -> None:
|
|
lib_paths = []
|
|
torch_root = _add_pytorch_path(pytorch_path)
|
|
|
|
if torch_root:
|
|
torch_lib_path = os.path.join(torch_root, "lib")
|
|
if os.path.isdir(torch_lib_path):
|
|
lib_paths.append(torch_lib_path)
|
|
|
|
if nvidia_path:
|
|
nvidia_root = os.path.abspath(nvidia_path)
|
|
nvjitlink_pattern = os.path.join(
|
|
nvidia_root, "nvjitlink", "lib", "libnvJitLink.so.*"
|
|
)
|
|
for library in [
|
|
"cublas",
|
|
"cudnn",
|
|
"cusparse",
|
|
"cusparse_full",
|
|
"nccl",
|
|
"nvshmem",
|
|
"nvjitlink",
|
|
]:
|
|
candidate = os.path.join(nvidia_root, library, "lib")
|
|
if os.path.isdir(candidate):
|
|
lib_paths.append(candidate)
|
|
|
|
if sys.platform.startswith("linux"):
|
|
nvjitlink_candidates = sorted(glob.glob(nvjitlink_pattern), reverse=True)
|
|
if nvjitlink_candidates:
|
|
ctypes.CDLL(nvjitlink_candidates[0], mode=ctypes.RTLD_GLOBAL)
|
|
|
|
_set_library_paths(lib_paths)
|
|
|
|
|
|
def _parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="Leon TCP server")
|
|
parser.add_argument(
|
|
"lang", nargs="?", default="en", help="Language code (e.g. en, fr)"
|
|
)
|
|
parser.add_argument("--pytorch-path", dest="pytorch_path", type=str, default=None)
|
|
parser.add_argument("--nvidia-path", dest="nvidia_path", type=str, default=None)
|
|
return parser.parse_args()
|
|
|
|
|
|
args = _parse_args()
|
|
os.environ["LEON_PY_TCP_SERVER_LANG"] = args.lang
|
|
_configure_external_libraries(args.pytorch_path, args.nvidia_path)
|
|
|
|
dotenv_path = join(resolve_leon_profile_path(), ".env")
|
|
load_dotenv(dotenv_path)
|
|
|
|
from lib.tcp_server import TCPServer
|
|
|
|
tcp_server_host = os.environ.get("LEON_PY_TCP_SERVER_HOST", DEFAULT_TCP_SERVER_HOST)
|
|
tcp_server_port = os.environ.get("LEON_PY_TCP_SERVER_PORT", DEFAULT_TCP_SERVER_PORT)
|
|
|
|
tcp_server = TCPServer(tcp_server_host, tcp_server_port)
|
|
|
|
# Use thread as ASR starts recording audio and it blocks the main thread
|
|
asr_thread = threading.Thread(target=tcp_server.init_asr)
|
|
asr_thread.start()
|
|
|
|
tcp_server.init_tts()
|
|
|
|
tcp_server_thread = threading.Thread(target=tcp_server.init)
|
|
tcp_server_thread.start()
|