|
|
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) |
|
|
|
|
|
|
