Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 09:14:35 +02:00
Code Issues Packages Projects Releases Wiki Activity

Don't assume Cuda on devices.same_device()

Browse Source
This commit is contained in:
Disty0
2024-10-14 17:23:51 +03:00
parent 31510cf5e4
commit b14e8f9a5f
2 changed files with 11 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/devices.py
+9 -5
View File
@@ -507,11 +507,15 @@ def test_for_nans(x, where):
raise NansException(message)
def normalize_device(device):
if torch.device(device).type in {"cpu", "mps", "meta"}:
return torch.device(device)
if torch.device(device).index is None:
return torch.device(str(device), index=0)
return torch.device(device)
def same_device(d1, d2):
if d1.type != d2.type:
return False
if d1.type == "cuda" and d1.index is None:
d1 = torch.device("cuda", index=0)
if d2.type == "cuda" and d2.index is None:
d2 = torch.device("cuda", index=0)
return d1 == d2
return normalize_device(d1) == normalize_device(d2)
modules/sd_models.py
+2 -10
View File
@@ -803,7 +803,7 @@ def apply_balanced_offload(sd_model):
return module
def pre_forward(self, module, *args, **kwargs):
if normalize_device(module.device) != normalize_device(devices.device):
if devices.normalize_device(module.device) != devices.normalize_device(devices.device):
device_index = torch.device(devices.device).index
if device_index is None:
device_index = 0
@@ -853,14 +853,6 @@ def apply_balanced_offload(sd_model):
return sd_model
def normalize_device(device):
if torch.device(device).type in {"cpu", "mps", "meta"}:
return torch.device(device)
if torch.device(device).index is None:
return torch.device(str(device) + ":0")
return torch.device(device)
def move_model(model, device=None, force=False):
if model is None or device is None:
return
@@ -900,7 +892,7 @@ def move_model(model, device=None, force=False):
shared.log.error(f'Model move execution device: device={device} {e}')
if getattr(model, 'has_accelerate', False) and not force:
return
if hasattr(model, "device") and normalize_device(model.device) == normalize_device(device):
if hasattr(model, "device") and devices.normalize_device(model.device) == devices.normalize_device(device):
return
try:
try: