mrq/tortoise-tts
4
6
Fork 15
Code Issues 9 Pull Requests 1 Packages Projects Releases Wiki Activity

expose VRAM easily

Browse Source
This commit is contained in:
mrq 2023-03-09 00:38:31 +00:00
parent 3dd5cad324
commit 6410df569b
1 changed files with 7 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
tortoise/utils/device.py
View File

@ -47,19 +47,19 @@ def get_device(verbose=False):
return torch.device(name) return torch.device(name)
def get_device_batch_size(): def get_device_vram( name=get_device_name() ):
available = 1 available = 1
name = get_device_name()
if name == "dml": if name == "cuda":
# there's nothing publically accessible in the DML API that exposes this
# there's a method to get currently used RAM statistics... as tiles
available = 1
elif name == "cuda":
_, available = torch.cuda.mem_get_info() _, available = torch.cuda.mem_get_info()
elif name == "cpu": elif name == "cpu":
available = psutil.virtual_memory()[4] available = psutil.virtual_memory()[4]
return available
def get_device_batch_size(name=None):
available = get_device_vram(name)
vram = available / (1024 ** 3) vram = available / (1024 ** 3)
# I'll need to rework this better # I'll need to rework this better
# simply adding more tiers clearly is not a good way to go about it # simply adding more tiers clearly is not a good way to go about it