From e7475ec69bba56358f99e8feffa5ebc45b957526 Mon Sep 17 00:00:00 2001
From: smthemex <138738845+smthemex@users.noreply.github.com>
Date: Wed, 5 Jun 2024 12:01:54 +0800
Subject: [PATCH] Update load_models_utils.py

add
---
 utils/load_models_utils.py | 25 +++++++++++++++++++------
 1 file changed, 19 insertions(+), 6 deletions(-)

diff --git a/utils/load_models_utils.py b/utils/load_models_utils.py
index 1835c91..3c61c52 100644
--- a/utils/load_models_utils.py
+++ b/utils/load_models_utils.py
@@ -1,6 +1,6 @@
 import yaml
 import torch
-from diffusers import StableDiffusionXLPipeline
+from diffusers import StableDiffusionXLPipeline, DiffusionPipeline
 from .pipeline import PhotoMakerStableDiffusionXLPipeline
 import os
 import sys
@@ -26,14 +26,14 @@ def get_models_dict():
             data = yaml.safe_load(stream)
             
             # 此时 'data' 是一个Python字典，里面包含了YAML文件的所有数据
-            #print(data)
+            # print(data)
             return data
             
         except yaml.YAMLError as exc:
             # 如果在解析过程中发生了错误，打印异常信息
             print(exc)
 
-def load_models(model_info,device,photomaker_path):
+def load_models(model_info,_sd_type,device,photomaker_path):
     path =  model_info["path"]
     single_files =  model_info["single_files"]
     use_safetensors = model_info["use_safetensors"]
@@ -46,7 +46,14 @@ def load_models(model_info,device,photomaker_path):
                 original_config_file=original_config_file,torch_dtype=torch.float16,
             )
         else:
-            pipe = StableDiffusionXLPipeline.from_pretrained(path, torch_dtype=torch.float16,)
+            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 = pipe.to(device)
     elif model_type == "Photomaker":
         if single_files:
@@ -56,8 +63,14 @@ def load_models(model_info,device,photomaker_path):
                 torch_dtype=torch.float16,use_safetensors=use_safetensors
             )
         else:
-            pipe = PhotoMakerStableDiffusionXLPipeline.from_pretrained(
-                path, torch_dtype=torch.float16, use_safetensors=use_safetensors)
+            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 = pipe.to(device)
         pipe.load_photomaker_adapter(
             os.path.dirname(photomaker_path),