# https://github.com/huggingface/huggingface_hub/blob/5a12851f54bf614be39614034ed3a9031922d297/src/huggingface_hub/utils/_runtime.py import os import platform import sys from pathlib import Path import packaging.version from .schema import Device from loguru import logger from rich import print from typing import Dict, Any # _PY_VERSION: str = sys.version.split()[0].rstrip("+") # if packaging.version.Version(_PY_VERSION) < packaging.version.Version("3.8.0"): # import importlib_metadata # type: ignore # else: # import importlib.metadata as importlib_metadata # type: ignore # _package_versions = {} # _CANDIDATES = [ # "torch", # "torchvision", # "Pillow", # "diffusers", # "transformers", # "opencv-python", # "accelerate", # "iopaint", # "rembg", # "realesrgan", # "gfpgan", # ] # Check once at runtime # for name in _CANDIDATES: # _package_versions[name] = "N/A" # try: # _package_versions[name] = importlib_metadata.version(name) # except importlib_metadata.PackageNotFoundError: # pass def dump_environment_info() -> Dict[str, str]: """Dump information about the machine to help debugging issues.""" # Generic machine info info: Dict[str, Any] = { "Platform": platform.platform(), "Python version": platform.python_version(), } info.update(_package_versions) print("\n".join([f"- {prop}: {val}" for prop, val in info.items()]) + "\n") return info def check_device(device: Device) -> Device: if device == Device.cuda: import platform if platform.system() == "Darwin": logger.warning("MacOS does not support cuda, use cpu instead") return Device.cpu else: import torch if not torch.cuda.is_available(): logger.warning("CUDA is not available, use cpu instead") return Device.cpu elif device == Device.mps: import torch if not torch.backends.mps.is_available(): logger.warning("mps is not available, use cpu instead") return Device.cpu return device def setup_model_dir(model_dir: Path): model_dir = model_dir.expanduser().absolute() logger.info(f"Model directory: {model_dir}") os.environ["U2NET_HOME"] = str(model_dir) os.environ["XDG_CACHE_HOME"] = str(model_dir) if not model_dir.exists(): logger.info(f"Create model directory: {model_dir}") model_dir.mkdir(exist_ok=True, parents=True) return model_dir