manager order

cloneofsimo · Feb 9, 2023 · 7be7bed · 7be7bed
1 parent 0ddce0b
commit 7be7bed
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 3 deletions.
diff --git a/example_loras/and.safetensors b/example_loras/and.safetensors
diff --git a/lora_diffusion/__init__.py b/lora_diffusion/__init__.py
@@ -2,3 +2,4 @@
 from .dataset import *
 from .utils import *
 from .preprocess_files import *
+from .lora_manager import *
diff --git a/lora_diffusion/lora_manager.py b/lora_diffusion/lora_manager.py
@@ -1,3 +1,4 @@
+from typing import List
 import torch
 from safetensors import safe_open
 from diffusers import StableDiffusionPipeline
@@ -50,12 +51,15 @@ def lora_join(lora_safetenors: list):
     for idx, safelora in enumerate(lora_safetenors):
         tokens = [k for k, v in safelora.metadata().items() if v == "<embed>"]
         for jdx, token in enumerate(sorted(tokens)):
-            if total_metadata.get(token, None) is not None:
-                del total_metadata[token]
+
             total_tensor[f"<s{idx}-{jdx}>"] = safelora.get_tensor(token)
             total_metadata[f"<s{idx}-{jdx}>"] = "<embed>"
+
             print(f"Embedding {token} replaced to <s{idx}-{jdx}>")
 
+            if total_metadata.get(token, None) is not None:
+                del total_metadata[token]
+
         token_size_list.append(len(tokens))
 
     return total_tensor, total_metadata, ranklist, token_size_list
@@ -77,7 +81,7 @@ def get_tensor(self, key):
 
 
 class LoRAManager:
-    def __init__(self, lora_paths_list, pipe: StableDiffusionPipeline):
+    def __init__(self, lora_paths_list: List[str], pipe: StableDiffusionPipeline):
 
         self.lora_paths_list = lora_paths_list
         self.pipe = pipe