104 lines
4.7 KiB
Python
104 lines
4.7 KiB
Python
import comfy.sd
|
|
import comfy.utils
|
|
import comfy.clip_vision
|
|
import comfy.model_detection
|
|
import comfy.model_management
|
|
import comfy.model_patcher
|
|
import comfy.checkpoint_pickle
|
|
|
|
import folder_paths
|
|
|
|
from .utils import download_file, load_torch_bin
|
|
|
|
class LoadTempMultiLoRA:
|
|
CKPT_TYPE = ["auto", "safetensors", "other"]
|
|
def __init__(self):
|
|
self.loaded_lora = {}
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"model": ("MODEL",),
|
|
"clip": ("CLIP", ),
|
|
"ckpt_url": ("STRING", {"default": "", "multiline": True, "dynamicPrompts": False}),
|
|
"ckpt_type": (LoadTempMultiLoRA.CKPT_TYPE, {"default": "auto"}),
|
|
"download_split": ("INT", {"default": 4, "min": 1, "max": 8, "step": 1}),
|
|
"strength_model": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}),
|
|
"strength_clip": ("FLOAT", {"default": 1.0, "min": -20.0, "max": 20.0, "step": 0.01}),
|
|
},
|
|
}
|
|
RETURN_TYPES = ("MODEL", "CLIP")
|
|
FUNCTION = "load_multi_lora"
|
|
|
|
CATEGORY = "temporary_loaders"
|
|
|
|
def load_multi_lora(self, model, clip, ckpt_url, ckpt_type, download_split, strength_model, strength_clip):
|
|
model_lora, clip_lora = model, clip
|
|
lora_dic = {}
|
|
for line in ckpt_url.splitlines():
|
|
if line.startswith("#"):
|
|
continue
|
|
else:
|
|
info = line.split(":")
|
|
if len(info) < 2:
|
|
continue
|
|
if line.startswith("file:"):
|
|
file_name = line[5:]
|
|
model_lora, clip_lora = self.load_lora_file(model_lora, clip_lora, file_name, strength_model, strength_clip, lora_dic)
|
|
continue
|
|
if line.startswith("http:") or line.startswith("https:"):
|
|
url = "{}:{}".format(info[0].strip(), info[1].strip())
|
|
model_lora, clip_lora = self.load_lora_url(model_lora, clip_lora, url, ckpt_type, download_split, strength_model, strength_clip, lora_dic)
|
|
continue
|
|
try:
|
|
isFile = info[-2].strip().startswith("file")
|
|
url = info[-1].strip() if isFile else "{}:{}".format(info[-2].strip(), info[-1].strip())
|
|
strength_m = float(info[0]) if len(info[0]) > 0 else strength_model
|
|
strength_c = float(info[1]) if len(info[1]) > 0 else strength_clip
|
|
m_type = info[2].strip() if len(info) == 5 and info[2].strip() in LoadTempMultiLoRA.CKPT_TYPE else ckpt_type
|
|
except:
|
|
raise Exception("Not according to format. The only accepted formats are\r\n{url}\r\n{strength_model}:{strength_clip}:{url}\r\n{strength_model}:{strength_clip}:{ckpt_type}:{url}")
|
|
if isFile:
|
|
model_lora, clip_lora = self.load_lora_file(model_lora, clip_lora, url, strength_m, strength_c, lora_dic)
|
|
else:
|
|
model_lora, clip_lora = self.load_lora_url(model_lora, clip_lora, url, m_type, download_split, strength_m, strength_c, lora_dic)
|
|
|
|
tmp = self.loaded_lora
|
|
self.loaded_lora = lora_dic
|
|
tmp.clear()
|
|
del tmp
|
|
|
|
return (model_lora, clip_lora)
|
|
|
|
def load_lora_url(self, model, clip, ckpt_url, ckpt_type, download_split, strength_model, strength_clip, lora_dic):
|
|
if strength_model == 0 and strength_clip == 0:
|
|
return (model, clip)
|
|
|
|
lora = self.loaded_lora.get(ckpt_url)
|
|
if lora is None:
|
|
bin, file_name = download_file(ckpt_url, download_split)
|
|
if bin is None:
|
|
raise file_name if file_name is not None else Exception("Download failed.")
|
|
|
|
is_safetensors = file_name.endswith(".safetensors") if ckpt_type =="auto" else ckpt_type == "safetensors"
|
|
lora = load_torch_bin(bin, is_safetensors, safe_load=True)
|
|
|
|
lora_dic[ckpt_url] = lora
|
|
|
|
model_lora, clip_lora = comfy.sd.load_lora_for_models(model, clip, lora, strength_model, strength_clip)
|
|
return (model_lora, clip_lora)
|
|
|
|
def load_lora_file(self, model, clip, lora_name, strength_model, strength_clip, lora_dic):
|
|
if strength_model == 0 and strength_clip == 0:
|
|
return (model, clip)
|
|
|
|
lora_path = folder_paths.get_full_path("loras", lora_name)
|
|
lora = self.loaded_lora.get(lora_path)
|
|
if lora is None:
|
|
lora = comfy.utils.load_torch_file(lora_path, safe_load=True)
|
|
lora_dic[lora_path] = lora
|
|
|
|
model_lora, clip_lora = comfy.sd.load_lora_for_models(model, clip, lora, strength_model, strength_clip)
|
|
return (model_lora, clip_lora)
|