81 lines
3.6 KiB
Python
81 lines
3.6 KiB
Python
"""Probe the PyTorch / CUDA situation, so the UI can show a device badge.
|
||
|
||
Pure (no Qt). ``gather()`` imports torch (slow / heavy) — call it off the GUI
|
||
thread. The rest are tiny formatters the UI uses to explain *why* it's on CPU and
|
||
how to enable the GPU.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
# pip index for the CUDA build (matches the README).
|
||
CUDA_WHEEL_INDEX = "https://download.pytorch.org/whl/cu121"
|
||
|
||
|
||
def gather() -> dict:
|
||
"""Collect torch/CUDA facts. Never raises — missing torch is a valid result."""
|
||
info: dict = {
|
||
"installed": False,
|
||
"version": None, # torch.__version__ (e.g. "2.12.0+cpu")
|
||
"built_cuda": None, # torch.version.cuda (None on a CPU-only build)
|
||
"cuda_available": False,
|
||
"device_name": None, # the active GPU's name, if any
|
||
"import_error": None,
|
||
}
|
||
try:
|
||
import torch
|
||
except Exception as exc: # noqa: BLE001 - report any import failure, not just ImportError
|
||
info["import_error"] = str(exc)
|
||
return info
|
||
info["installed"] = True
|
||
info["version"] = getattr(torch, "__version__", None)
|
||
try:
|
||
info["built_cuda"] = torch.version.cuda
|
||
except Exception: # noqa: BLE001
|
||
info["built_cuda"] = None
|
||
try:
|
||
info["cuda_available"] = bool(torch.cuda.is_available())
|
||
except Exception: # noqa: BLE001
|
||
info["cuda_available"] = False
|
||
if info["cuda_available"]:
|
||
try:
|
||
info["device_name"] = torch.cuda.get_device_name(0)
|
||
except Exception: # noqa: BLE001
|
||
info["device_name"] = None
|
||
return info
|
||
|
||
|
||
def device_label(info: dict) -> str:
|
||
return "CUDA" if info.get("cuda_available") else "CPU"
|
||
|
||
|
||
def reason(info: dict) -> str:
|
||
"""One-sentence human explanation of the current device choice."""
|
||
if not info.get("installed"):
|
||
return ("PyTorch не установлен — детектор YOLO и восстановление DeepMosaics "
|
||
"работают на CPU (классический детектор torch не требует).")
|
||
if info.get("cuda_available"):
|
||
name = info.get("device_name") or "GPU"
|
||
return f"PyTorch использует CUDA: {name}. Вычисления идут на видеокарте."
|
||
version = info.get("version") or "?"
|
||
built = info.get("built_cuda")
|
||
if not built:
|
||
return (f"Установлена CPU-сборка PyTorch ({version}) — без поддержки CUDA, "
|
||
"поэтому вычисления идут на процессоре (медленно).")
|
||
return (f"PyTorch собран с CUDA {built} ({version}), но GPU недоступен: нет "
|
||
"NVIDIA-видеокарты, не установлен/устарел драйвер, либо версия CUDA "
|
||
"несовместима с драйвером.")
|
||
|
||
|
||
def install_hint() -> str:
|
||
"""Steps to enable the GPU (shown when running on CPU)."""
|
||
return (
|
||
"Как включить GPU (NVIDIA):\n"
|
||
"1. Нужна видеокарта NVIDIA и свежий драйвер (проверка в консоли: nvidia-smi).\n"
|
||
"2. Переустановите PyTorch со сборкой CUDA:\n\n"
|
||
" pip uninstall -y torch torchvision\n"
|
||
f" pip install torch torchvision --index-url {CUDA_WHEEL_INDEX}\n\n"
|
||
"3. Перезапустите приложение.\n\n"
|
||
"Классический детектор работает и без CUDA. На CPU детекция и расцензуривание "
|
||
"просто медленнее."
|
||
)
|