nodes/smthemex-ComfyUI_StoryDiffusion
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update load_models_utils.py

Browse Source 
add
This commit is contained in:
smthemex
2024-06-16 11:59:41 +08:00
committed by GitHub
parent d86ae2bfc5
commit 225ce91bc0
1 changed files with 33 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
utils/load_models_utils.py
+33 -10
View File
@@ -4,6 +4,9 @@ from diffusers import StableDiffusionXLPipeline, DiffusionPipeline
from .pipeline import PhotoMakerStableDiffusionXLPipeline
import os
import sys
import diffusers
dif_version = str(diffusers.__version__)
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)
@@ -33,7 +36,7 @@ def get_models_dict():
# 如果在解析过程中发生了错误，打印异常信息
print(exc)
def load_models(model_info,_sd_type,device,photomaker_path):
def load_models(model_info,_sd_type,device,photomaker_path,lora,trigger_words):
path = model_info["path"]
single_files = model_info["single_files"]
use_safetensors = model_info["use_safetensors"]
@@ -41,10 +44,15 @@ def load_models(model_info,_sd_type,device,photomaker_path):
if model_type == "original":
if single_files:
pipe = StableDiffusionXLPipeline.from_single_file(
pretrained_model_link_or_path=path,
original_config_file=original_config_file,torch_dtype=torch.float16,
)
if dif_version_int>=28:
pipe = StableDiffusionXLPipeline.from_single_file(
pretrained_model_link_or_path=path, original_config=original_config_file, torch_dtype=torch.float16)
else:
pipe = StableDiffusionXLPipeline.from_single_file(
pretrained_model_link_or_path=path, original_config_file=original_config_file, torch_dtype=torch.float16
)
else:
if _sd_type=="Playground_v2p5":
pipe = DiffusionPipeline.from_pretrained(
@@ -54,14 +62,24 @@ def load_models(model_info,_sd_type,device,photomaker_path):
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.set_adapters(trigger_words)
pipe.fuse_lora()
# pipe._lora_scale=lora_scale
elif model_type == "Photomaker":
if single_files:
print("loading from a single_files")
pipe = PhotoMakerStableDiffusionXLPipeline.from_single_file(
pretrained_model_link_or_path=path, original_config_file=original_config_file,
torch_dtype=torch.float16,use_safetensors=use_safetensors
)
#print("loading from a single_files")
if dif_version_int>=28:
pipe = PhotoMakerStableDiffusionXLPipeline.from_single_file(
pretrained_model_link_or_path=path, original_config=original_config_file, torch_dtype=torch.float16,use_safetensors=use_safetensors)
else:
pipe = PhotoMakerStableDiffusionXLPipeline.from_single_file(
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(
@@ -71,6 +89,7 @@ def load_models(model_info,_sd_type,device,photomaker_path):
pipe = PhotoMakerStableDiffusionXLPipeline.from_pretrained(
path, torch_dtype=torch.float16,use_safetensors=use_safetensors
)
pipe = pipe.to(device)
pipe.load_photomaker_adapter(
os.path.dirname(photomaker_path),
@@ -78,6 +97,10 @@ def load_models(model_info,_sd_type,device,photomaker_path):
weight_name=os.path.basename(photomaker_path),
trigger_word="img" # define the trigger word
)
if lora != "none":
pipe.load_lora_weights(lora, adapter_name=trigger_words)
pipe.set_adapters(trigger_words)
# pipe._lora_scale=lora_scale
pipe.fuse_lora()
else:
raise NotImplementedError("You should choice between original and Photomaker!",f"But you choice {model_type}")