Fix torch.linalg.solve with IPEX & Diffusers UniPC

This commit is contained in:
Disty0
2023-06-10 22:01:09 +03:00
parent baecfb7a13
commit f63dd1c92e
2 changed files with 6 additions and 15 deletions
+4 -3
View File
@@ -34,9 +34,7 @@ def get_cuda_device_string():
def get_optimal_device_name():
if shared.cmd_opts.use_ipex:
return get_cuda_device_string()
elif cuda_ok and not shared.cmd_opts.use_directml:
if (cuda_ok or shared.cmd_opts.use_ipex) and not shared.cmd_opts.use_directml:
return get_cuda_device_string()
if has_mps():
return "mps"
@@ -172,6 +170,9 @@ if args.use_ipex:
CondFunc('torch.nn.modules.Linear.forward',
lambda orig_func, *args, **kwargs: orig_func(args[0], args[1].to(args[0].weight.data.dtype)),
lambda *args, **kwargs: args[2].dtype != args[1].weight.data.dtype)
CondFunc('torch.linalg.solve',
lambda orig_func, *args, **kwargs: orig_func(args[0].to("cpu"), args[1].to("cpu")).to(get_cuda_device_string()),
lambda *args, **kwargs: True)
#Use XPU instead of CPU. %20 Perf improvement on weak CPUs.
if args.device_id is not None:
+2 -12
View File
@@ -685,12 +685,7 @@ class UniPC:
if order == 2:
rhos_p = torch.tensor([0.5], device=b.device)
else:
if shared.cmd_opts.use_ipex:
#Running torch.linalg.solve on XPU crashes the GPU.
rhos_p = torch.linalg.solve(R[:-1, :-1].to("cpu"), b[:-1].to("cpu"))
rhos_p = rhos_p.to(b.device)
else:
rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1])
rhos_p = torch.linalg.solve(R[:-1, :-1], b[:-1])
else:
D1s = None
@@ -700,12 +695,7 @@ class UniPC:
if order == 1:
rhos_c = torch.tensor([0.5], device=b.device)
else:
if shared.cmd_opts.use_ipex:
#Running torch.linalg.solve on XPU crashes the GPU.
rhos_c = torch.linalg.solve(R.to("cpu"), b.to("cpu"))
rhos_c = rhos_c.to(b.device)
else:
rhos_c = torch.linalg.solve(R, b)
rhos_c = torch.linalg.solve(R, b)
model_t = None
if self.predict_x0: