From 42c6147ac8d6c316a558e5fb1811b8b17a2bb8b3 Mon Sep 17 00:00:00 2001
From: Seunghoon Lee <lshqqytiger@naver.com>
Date: Sat, 29 Jul 2023 13:48:28 +0900
Subject: [PATCH] cleanup

---
 modules/dml/__init__.py | 6 +++---
 modules/dml/backend.py  | 8 ++++----
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/modules/dml/__init__.py b/modules/dml/__init__.py
index ff1782542..3c8d8f0e1 100644
--- a/modules/dml/__init__.py
+++ b/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."),
diff --git a/modules/dml/backend.py b/modules/dml/backend.py
index c7d9ac0ef..feb0a07bd 100644
--- a/modules/dml/backend.py
+++ b/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: