|
|
|
|
|
|
|
device = "cuda:0"
|
|
|
|
|
|
is_half = True
|
|
|
|
|
|
n_cpu = 0
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import argparse
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--port", type=int, default=7865, help="Listen port")
|
|
parser.add_argument("--pycmd", type=str, default="python", help="Python command")
|
|
parser.add_argument("--colab", action="store_true", help="Launch in colab")
|
|
parser.add_argument(
|
|
"--noparallel", action="store_true", help="Disable parallel processing"
|
|
)
|
|
parser.add_argument(
|
|
"--noautoopen", action="store_true", help="Do not open in browser automatically"
|
|
)
|
|
cmd_opts = parser.parse_args()
|
|
|
|
python_cmd = cmd_opts.pycmd
|
|
listen_port = cmd_opts.port if 0 <= cmd_opts.port <= 65535 else 7865
|
|
iscolab = cmd_opts.colab
|
|
noparallel = cmd_opts.noparallel
|
|
noautoopen = cmd_opts.noautoopen
|
|
|
|
|
|
import sys
|
|
import torch
|
|
|
|
|
|
|
|
|
|
def has_mps() -> bool:
|
|
if sys.platform != "darwin":
|
|
return False
|
|
else:
|
|
if not getattr(torch, "has_mps", False):
|
|
return False
|
|
try:
|
|
torch.zeros(1).to(torch.device("mps"))
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
if not torch.cuda.is_available():
|
|
if has_mps():
|
|
print("没有发现支持的N卡, 使用MPS进行推理")
|
|
device = "mps"
|
|
else:
|
|
print("没有发现支持的N卡, 使用CPU进行推理")
|
|
device = "cpu"
|
|
is_half = False
|
|
|
|
if device not in ["cpu", "mps"]:
|
|
gpu_name = torch.cuda.get_device_name(int(device.split(":")[-1]))
|
|
if "16" in gpu_name or "MX" in gpu_name:
|
|
print("16系显卡/MX系显卡强制单精度")
|
|
is_half = False
|
|
|
|
from multiprocessing import cpu_count
|
|
|
|
if n_cpu == 0:
|
|
n_cpu = cpu_count()
|
|
if is_half:
|
|
|
|
x_pad = 3
|
|
x_query = 10
|
|
x_center = 60
|
|
x_max = 65
|
|
else:
|
|
|
|
x_pad = 1
|
|
x_query = 6
|
|
x_center = 38
|
|
x_max = 41
|
|
|