From b5d1af11b7dc718d4d91d379c75e46f4bd2e2fe6 Mon Sep 17 00:00:00 2001 From: Abdullah Barhoum Date: Sun, 11 Sep 2022 07:11:27 +0200 Subject: Modular device management --- modules/lowvram.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) (limited to 'modules/lowvram.py') diff --git a/modules/lowvram.py b/modules/lowvram.py index bd117491..079386c3 100644 --- a/modules/lowvram.py +++ b/modules/lowvram.py @@ -1,13 +1,9 @@ import torch +from modules.devices import get_optimal_device module_in_gpu = None cpu = torch.device("cpu") -if torch.has_cuda: - device = gpu = torch.device("cuda") -elif torch.has_mps: - device = gpu = torch.device("mps") -else: - device = gpu = torch.device("cpu") +device = gpu = get_optimal_device() def setup_for_low_vram(sd_model, use_medvram): parents = {} -- cgit v1.2.1