set[type] also doesnt work on older Pythons, replace with the more co…

…mpatible Set[type]

lora_diffusion/lora.py

@@ -1,6 +1,6 @@
 import math
 from itertools import groupby
-from typing import Callable, Dict, List, Optional, Tuple, Union
+from typing import Callable, Dict, List, Optional, Set, Tuple, Union
 
 import numpy as np
 import PIL
@@ -55,7 +55,7 @@ def _find_children(
 
 def _find_modules(
     model,
-    ancestor_class: set[str] = DEFAULT_TARGET_REPLACE,
+    ancestor_class: Set[str] = DEFAULT_TARGET_REPLACE,
     search_class: list[type[nn.Module]] = [nn.Linear],
     exclude_children_of: Optional[list[type[nn.Module]]] = [LoraInjectedLinear],
 ):
@@ -94,7 +94,7 @@ def _find_modules(
 
 def inject_trainable_lora(
     model: nn.Module,
-    target_replace_module: set[str] = DEFAULT_TARGET_REPLACE,
+    target_replace_module: Set[str] = DEFAULT_TARGET_REPLACE,
     r: int = 4,
     loras=None,  # path to lora .pt
 ):
@@ -183,7 +183,7 @@ def save_lora_as_json(model, path="./lora.json"):
 
 
 def save_safeloras(
-    modelmap: dict[str, tuple[nn.Module, set[str]]] = {},
+    modelmap: dict[str, tuple[nn.Module, Set[str]]] = {},
     outpath="./lora.safetensors",
 ):
     """
@@ -215,7 +215,7 @@ def save_safeloras(
 
 
 def convert_loras_to_safeloras(
-    modelmap: dict[str, tuple[str, set[str], int]] = {},
+    modelmap: dict[str, tuple[str, Set[str], int]] = {},
     outpath="./lora.safetensors",
 ):
     """