File size: 1,120 Bytes
d4d21ad | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 | import importlib.util
import torch
from .sample_funcs import sample_naive
from .structs import T2SRequest, T2SResult
from .t2s_engine import T2SEngine as T2SEngineTorch
torch.set_grad_enabled(False)
backends = ["torch_varlen"]
if torch.cuda.is_available():
backends.append("torch_static_cuda_graph")
# if importlib.util.find_spec("sageattention") is not None:
# for i in range(torch.cuda.device_count()):
# major, minor = torch.cuda.get_device_capability(i)
# sm_version = major + minor / 10.0
# if sm_version >= 7.0:
# backends.append("sage_attn_varlen_cuda_graph")
if importlib.util.find_spec("flash_attn") is not None:
for i in range(torch.cuda.device_count()):
major, minor = torch.cuda.get_device_capability(i)
sm_version = major + minor / 10.0
if sm_version >= 7.5:
backends.append("flash_attn_varlen_cuda_graph")
# if torch.mps.is_available():
# backends.append("mps_flash_attn_varlen")
__all__ = ["T2SEngineTorch", "T2SRequest", "sample_naive", "T2SResult", "backends"]
|