(requested) support AMD 8GB GPUs via Windows DirectML

this update is requested by users
This commit is contained in:
lllyasviel
2023-12-30 06:30:59 -08:00
committed by GitHub
parent c0e11c3451
commit 8e62a72a63
31 changed files with 402 additions and 148 deletions
+19
View File
@@ -0,0 +1,19 @@
import torch
import contextlib
@contextlib.contextmanager
def use_patched_ops(operations):
op_names = ['Linear', 'Conv2d', 'Conv3d', 'GroupNorm', 'LayerNorm']
backups = {op_name: getattr(torch.nn, op_name) for op_name in op_names}
try:
for op_name in op_names:
setattr(torch.nn, op_name, getattr(operations, op_name))
yield
finally:
for op_name in op_names:
setattr(torch.nn, op_name, backups[op_name])
return