From 11205a8a306038b25e496052dfb09d10140e1390 Mon Sep 17 00:00:00 2001
From: smthemex <138738845+smthemex@users.noreply.github.com>
Date: Sat, 22 Jun 2024 20:45:22 +0800
Subject: [PATCH] Update load_models_utils.py

update
---
 utils/load_models_utils.py | 73 +++++++++++++++++++++++---------------
 1 file changed, 45 insertions(+), 28 deletions(-)

diff --git a/utils/load_models_utils.py b/utils/load_models_utils.py
index f5592f1..0c0ce8e 100644
--- a/utils/load_models_utils.py
+++ b/utils/load_models_utils.py
@@ -10,20 +10,38 @@ dif_version_int= int(dif_version.split(".")[1])
 dir_path = os.path.dirname(os.path.abspath(__file__))
 path_dir = os.path.dirname(dir_path)
 file_path = os.path.dirname(path_dir)
-yaml_path = os.path.join(path_dir ,'config','models.yaml')
-original_config_file=os.path.join(path_dir,'config','sd_xl_base.yaml')
+
 def get_instance_path(path):
     instance_path = os.path.normpath(path)
     if sys.platform == 'win32':
         instance_path = path.replace('\\', "/")
     return instance_path
 
-yaml_path=get_instance_path(yaml_path)
-original_config_file = get_instance_path(original_config_file)
-#print(yaml_path,original_config_file)
+yaml_path = get_instance_path(os.path.join(path_dir ,'config','models.yaml'))
+original_config_file=get_instance_path(os.path.join(path_dir,'config','sd_xl_base.yaml'))
+loras_path = get_instance_path(os.path.join(path_dir,"config","lora.yaml"))
+
+def get_lora_dict():
+    # 打开并读取YAML文件
+    with open(loras_path, 'r', encoding="UTF-8") as stream:
+        try:
+            # 解析YAML文件内容
+            data = yaml.safe_load(stream)
+
+            # 此时 'data' 是一个Python字典，里面包含了YAML文件的所有数据
+            # print(data)
+            return data
+
+        except yaml.YAMLError as exc:
+            # 如果在解析过程中发生了错误，打印异常信息
+            print(exc)
+
+datas = get_lora_dict()
+lora_lightning_list = datas["lightning_xl_lora"]
+#print(lora_lightning_list)
 def get_models_dict():
     # 打开并读取YAML文件
-    with open(yaml_path, 'r') as stream:
+    with open(yaml_path, 'r', encoding="UTF-8") as stream:
         try:
             # 解析YAML文件内容
             data = yaml.safe_load(stream)
@@ -36,7 +54,7 @@ def get_models_dict():
             # 如果在解析过程中发生了错误，打印异常信息
             print(exc)
 
-def load_models(model_info,_sd_type,device,photomaker_path,lora,trigger_words,lora_scale):
+def load_models(model_info,_sd_type,device,photomaker_path,lora,lora_path,trigger_words,lora_scale):
     path =  model_info["path"]
     single_files =  model_info["single_files"]
     use_safetensors = model_info["use_safetensors"]
@@ -54,19 +72,19 @@ def load_models(model_info,_sd_type,device,photomaker_path,lora,trigger_words,lo
                 )
 
         else:
-            if _sd_type=="Playground_v2p5":
-                pipe = DiffusionPipeline.from_pretrained(
-                    path, torch_dtype=torch.float16, use_safetensors=use_safetensors
-                )
-            else:
-                pipe = StableDiffusionXLPipeline.from_pretrained(
-                    path, torch_dtype=torch.float16,use_safetensors=use_safetensors
-                )
+            pipe = StableDiffusionXLPipeline.from_pretrained(
+                path, torch_dtype=torch.float16,use_safetensors=use_safetensors
+            )
 
         pipe = pipe.to(device)
         if lora != "none":
-            pipe.load_lora_weights(lora, adapter_name=trigger_words)
-            pipe.fuse_lora(adapter_names=[trigger_words,], lora_scale=lora_scale)
+            if lora in lora_lightning_list:
+                pipe.load_lora_weights(lora_path)
+                pipe.fuse_lora()
+            else:
+                pipe.load_lora_weights(lora_path, adapter_name=trigger_words)
+                pipe.fuse_lora(adapter_names=[trigger_words, ], lora_scale=lora_scale)
+
     elif model_type == "Photomaker":
         if single_files:
             #print("loading from a single_files")
@@ -78,25 +96,24 @@ def load_models(model_info,_sd_type,device,photomaker_path,lora,trigger_words,lo
                     pretrained_model_link_or_path=path, original_config_file=original_config_file, torch_dtype=torch.float16,use_safetensors=use_safetensors
                 )
         else:
-            if _sd_type=="Playground_v2p5":
-                pipe = PhotoMakerStableDiffusionXLPipeline.from_pretrained(
-                    path, torch_dtype=torch.float16, use_safetensors=use_safetensors
-                )
-            else:
-                pipe = PhotoMakerStableDiffusionXLPipeline.from_pretrained(
-                    path, torch_dtype=torch.float16,use_safetensors=use_safetensors
-                )
+            pipe = PhotoMakerStableDiffusionXLPipeline.from_pretrained(
+                path, torch_dtype=torch.float16,use_safetensors=use_safetensors
+            )
 
         pipe = pipe.to(device)
-        if lora != "none":
-            pipe.load_lora_weights(lora, adapter_name=trigger_words)
-            pipe.fuse_lora(adapter_names=[trigger_words,], lora_scale=lora_scale)
         pipe.load_photomaker_adapter(
             os.path.dirname(photomaker_path),
             subfolder="",
             weight_name=os.path.basename(photomaker_path),
             trigger_word="img"  # define the trigger word
         )
+        if lora != "none":
+            if lora in lora_lightning_list:
+                pipe.load_lora_weights(lora_path)
+                pipe.fuse_lora()
+            else:
+                pipe.load_lora_weights(lora_path, adapter_name=trigger_words)
+                pipe.fuse_lora(adapter_names=[trigger_words, ], lora_scale=lora_scale)
 
 
     else: