61 lines
1.6 KiB
Python
61 lines
1.6 KiB
Python
import folder_paths
|
|
from comfy.utils import load_torch_file
|
|
|
|
class StateDictLoader:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
'required': {
|
|
'ckpt_name': (folder_paths.get_filename_list("checkpoints"), )
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ('DICT',)
|
|
|
|
FUNCTION = 'execute'
|
|
|
|
CATEGORY = 'loaders'
|
|
|
|
def execute(self, ckpt_name: str):
|
|
ckpt_path = folder_paths.get_full_path('checkpoints', ckpt_name)
|
|
sd = load_torch_file(ckpt_path)
|
|
return (sd,)
|
|
|
|
|
|
class Dict2Model:
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
'required': {
|
|
'weights': ('DICT', ),
|
|
'config_name': (folder_paths.get_filename_list('configs'), ),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ('MODEL', 'CLIP', 'VAE')
|
|
|
|
FUNCTION = 'execute'
|
|
|
|
CATEGORY = 'model'
|
|
|
|
def execute(self, weights: dict, config_name: str):
|
|
config_path = folder_paths.get_full_path("configs", config_name)
|
|
|
|
def load_torch_file_hook(*args, **kwargs):
|
|
return weights
|
|
|
|
from comfy import utils, sd
|
|
load_torch_file_org = utils.load_torch_file
|
|
setattr(utils, 'load_torch_file', load_torch_file_hook)
|
|
|
|
try:
|
|
model, clip, vae = sd.load_checkpoint(config_path, None, output_vae=True, output_clip=True, embedding_directory=folder_paths.get_folder_paths("embeddings"))
|
|
assert clip is not None
|
|
assert vae is not None
|
|
return (model, clip, vae)
|
|
finally:
|
|
setattr(sd, 'load_torch_file', load_torch_file_org)
|
|
|