cleanup

modules/dml/__init__.py

@@ -1,6 +1,6 @@
 import os
 import torch
-from typing import NamedTuple, Callable, Union
+from typing import NamedTuple, Callable, Optional
 
 from modules.sd_hijack_utils import CondFunc
 
@@ -44,8 +44,8 @@ def directml_do_hijack():
 
 class OverrideItem(NamedTuple):
     value: str
-    condition: Union[Callable, None]
-    message: Union[str, None]
+    condition: Optional[Callable]
+    message: Optional[str]
 
 opts_override_table = {
     "diffusers_generator_device": OverrideItem("cpu", None, "DirectML does not support torch Generator API."),

modules/dml/backend.py

@@ -1,5 +1,5 @@
 # pylint: disable=no-member,no-self-argument,no-method-argument
-from typing import Optional, Union
+from typing import Optional
 import torch
 import torch_directml # pylint: disable=import-error
 import modules.dml.amp as amp
@@ -13,9 +13,9 @@ class DirectML:
     amp = amp
     device = device
 
-    context_device: Union[torch.device, None] = None
+    context_device: Optional[torch.device] = None
 
-    __gpu_memory_bound: Union[int, None] = None
+    __gpu_memory_bound: Optional[int] = None
     
     is_autocast_enabled = False
     autocast_gpu_dtype = torch.float16
@@ -36,7 +36,7 @@ class DirectML:
         except Exception:
             return UnknownMemoryControl
         
-    def set_gpu_memory_bound(bound: Union[int, None]):
+    def set_gpu_memory_bound(bound: Optional[int]):
         DirectML.__gpu_memory_bound = bound
 
     def is_available() -> bool: