Mirrored_Repos/InvokeAI
1
0
Fork 0
mirror of https://github.com/invoke-ai/InvokeAI synced 2024-08-30 20:32:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

all vestiges of ldm.invoke removed

Browse Source
This commit is contained in:
Lincoln Stein
2023-03-03 01:02:00 -05:00
parent 6a990565ff
commit 60a98cacef
126 changed files with 8514 additions and 6520 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
invokeai/backend/util/devices.py
View File

@ -11,46 +11,51 @@ CPU_DEVICE = torch.device("cpu")
CUDA_DEVICE = torch.device("cuda")
MPS_DEVICE = torch.device("mps")
def choose_torch_device() -> torch.device:
'''Convenience routine for guessing which GPU device to run model on'''
"""Convenience routine for guessing which GPU device to run model on"""
if Globals.always_use_cpu:
return CPU_DEVICE
if torch.cuda.is_available():
return torch.device('cuda')
if hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
return torch.device('mps')
return torch.device("cuda")
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
return torch.device("mps")
return CPU_DEVICE
def choose_precision(device: torch.device) -> str:
'''Returns an appropriate precision for the given torch device'''
if device.type == 'cuda':
"""Returns an appropriate precision for the given torch device"""
if device.type == "cuda":
device_name = torch.cuda.get_device_name(device)
if not ('GeForce GTX 1660' in device_name or 'GeForce GTX 1650' in device_name):
return 'float16'
return 'float32'
if not ("GeForce GTX 1660" in device_name or "GeForce GTX 1650" in device_name):
return "float16"
return "float32"
def torch_dtype(device: torch.device) -> torch.dtype:
if Globals.full_precision:
return torch.float32
if choose_precision(device) == 'float16':
if choose_precision(device) == "float16":
return torch.float16
else:
return torch.float32
def choose_autocast(precision):
'''Returns an autocast context or nullcontext for the given precision string'''
"""Returns an autocast context or nullcontext for the given precision string"""
# float16 currently requires autocast to avoid errors like:
# 'expected scalar type Half but found Float'
if precision == 'autocast' or precision == 'float16':
if precision == "autocast" or precision == "float16":
return autocast
return nullcontext
def normalize_device(device: str | torch.device) -> torch.device:
"""Ensure device has a device index defined, if appropriate."""
device = torch.device(device)
if device.index is None:
# cuda might be the only torch backend that currently uses the device index?
# I don't see anything like `current_device` for cpu or mps.
if device.type == 'cuda':
if device.type == "cuda":
device = torch.device(device.type, torch.cuda.current_device())
return device