Combine selectors. Cleanup.

This commit is contained in:
arcum42
2025-08-03 01:29:13 -07:00
parent e469a795a2
commit fbecfa5c72
3 changed files with 81 additions and 49 deletions
+3 -5
View File
@@ -92,8 +92,7 @@ LOADER_CLASS_MAPPINGS = {
"Sage_CLIPLoaderFromInfo": Sage_CLIPLoaderFromInfo,
"Sage_VAELoaderFromInfo": Sage_VAELoaderFromInfo,
"Sage_LoraStackLoader": Sage_LoraStackLoader,
"Sage_ModelLoraStackLoader": Sage_ModelLoraStackLoader,
"Sage_GGUFTestUNETLoaderSimple": Sage_GGUFTestUNETLoaderSimple
"Sage_ModelLoraStackLoader": Sage_ModelLoraStackLoader
}
MODEL_CLASS_MAPPINGS = {
@@ -214,7 +213,7 @@ SELECTOR_NAME_MAPPINGS = {
"Sage_TripleCLIPSelector": "Triple CLIP Selector",
"Sage_QuadCLIPSelector": "Quad CLIP Selector",
"Sage_VAESelector": "VAE Selector",
"Sage_UnetClipVaeToModelInfo": "UNET + CLIP + VAE to Model Info"
"Sage_UnetClipVaeToModelInfo": "UNET + CLIP + VAE"
}
TEXT_NAME_MAPPINGS = {
@@ -254,8 +253,7 @@ MODEL_NAME_MAPPINGS = {
"Sage_CacheMaintenance": "Cache Maintenance",
"Sage_ModelReport": "Model Scan & Report",
"Sage_MultiModelPicker": "Multi Model Picker",
"Sage_GGUFTestUNETLoaderSimple": "GGUF Test UNET Loader",
"Sage_LoadModelwithGraphBuilderTest": "Load Model with Graph Builder Test"
"Sage_LoadModelwithGraphBuilderTest": "Load Models"
}
LORA_NAME_MAPPINGS = {
+57 -38
View File
@@ -417,6 +417,9 @@ def get_path_without_base(folder_type:str, path:str) -> str:
break
return path
def get_file_extension(path: str) -> str:
"""Get the file extension from a path."""
return path.split(".")[-1] if "." in path else ""
class Sage_LoadModelwithGraphBuilderTest(ComfyNodeABC):
@classmethod
def INPUT_TYPES(cls) -> InputTypeDict:
@@ -436,18 +439,13 @@ class Sage_LoadModelwithGraphBuilderTest(ComfyNodeABC):
CATEGORY = "Sage Utils/model"
DESCRIPTION = "Load model components from model info using GraphBuilder."
def load_model(self, model_info):
def prepare_model_graph(self, model_info):
graph = GraphBuilder()
nodes = []
ckpt_node = None
unet_node = None
clip_node = None
vae_node = None
ckpt_node = unet_node = clip_node = vae_node = None
clip_out = unet_out = vae_out = None
ckpt_info = unet_info = clip_info = vae_info = None
clip_out = None
unet_out = None
vae_out = None
# We have a few scenerios to handle.
# We need to loop over the model_info and see what types of models we have.
# We'll worry about GGUF later.
@@ -461,21 +459,6 @@ class Sage_LoadModelwithGraphBuilderTest(ComfyNodeABC):
# A checkpoint provides UNET, CLIP, and VAE. If we have the others, they will override the checkpoint components.
# Most likely scenerio is either a checkpoint, or UNET, CLIP, and VAE.
# If we have a checkpoint, we will load it first, then load the others.
ckpt_info = None
unet_info = None
clip_info = None
vae_info = None
# Loop through the model_info to determine what types of models we have, and record their indices.
# Don't use get_model_types here. Seriously, don't.
#print(f"Model info type is {type(model_info)}; items: {model_info}")
# Model info type is <class 'list'>; items: [{'type': 'CKPT', 'path': '/home/ai/models/checkpoints/ponyXL/ponyDiffusionV6XL_v6.safetensors', 'hash': 'f7564bb179'}]
# Checking model info item 0: {'type': 'CKPT', 'path': '/home/ai/models/checkpoints/ponyXL/ponyDiffusionV6XL_v6.safetensors', 'hash': 'f7564bb179'}
# Model info type is <class 'list'>; items: [({'type': 'UNET', 'path': '/home/ai/models/diffusion_models/chroma/chroma-unlocked-v48_float8_e4m3fn_scaled_learned_nodistill.safetensors', 'hash': '727416bfb5', 'weight_dtype': 'default'}, {'type': 'CLIP', 'path': ['/home/ai/models/clip/t5/t5xxl_fp8_e4m3fn_scaled.safetensors'], 'hash': ['a498f0485d'], 'clip_type': 'chroma'}, {'type': 'VAE', 'path': '/home/ai/models/vae/ae.safetensors', 'hash': 'afc8e28272'})]
# Checking model info item 0: ({'type': 'UNET', 'path': '/home/ai/models/diffusion_models/chroma/chroma-unlocked-v48_float8_e4m3fn_scaled_learned_nodistill.safetensors', 'hash': '727416bfb5', 'weight_dtype': 'default'}, {'type': 'CLIP', 'path': ['/home/ai/models/clip/t5/t5xxl_fp8_e4m3fn_scaled.safetensors'], 'hash': ['a498f0485d'], 'clip_type': 'chroma'}, {'type': 'VAE', 'path': '/home/ai/models/vae/ae.safetensors', 'hash': 'afc8e28272'})
# No checkpoint found in model info, skipping checkpoint loading.
# model_info is either a list of dicts, or a list of tuples with dicts inside. Either way, we need to go over each dict, and check the "type" key.
@@ -515,6 +498,7 @@ class Sage_LoadModelwithGraphBuilderTest(ComfyNodeABC):
if ckpt_info is not None:
ckpt_name = ckpt_info["path"]
ckpt_fixed_name = get_path_without_base("checkpoints", ckpt_name)
# No GGUF support for checkpoints yet that I see.
ckpt_node = graph.node("CheckpointLoaderSimple",ckpt_name=ckpt_fixed_name)
print(f"Checkpoint node created with name: {ckpt_fixed_name}")
else:
@@ -525,8 +509,17 @@ class Sage_LoadModelwithGraphBuilderTest(ComfyNodeABC):
unet_name = unet_info["path"]
unet_weight_dtype = unet_info["weight_dtype"]
unet_fixed_name = get_path_without_base("diffusion_models", unet_name)
unet_node = graph.node("UNETLoader", unet_name=unet_fixed_name, weight_dtype=unet_weight_dtype)
print(f"UNET node created with name: {unet_fixed_name}, weight_dtype: {unet_weight_dtype}")
if get_file_extension(unet_fixed_name) == "gguf":
# If the UNET is a GGUF, we need to use the GGUFLoader node. It doesn't ask for weight_dtype.
try:
unet_node = graph.node("UnetLoaderGGUF", unet_name=unet_fixed_name)
print(f"GGUF UNET node created with name: {unet_fixed_name}")
except:
print("Unable to load UNET as GGUF. Do you have ComfyUI-GGUF installed?")
raise ValueError("Unable to load UNET as GGUF. Do you have ComfyUI-GGUF installed?")
else:
unet_node = graph.node("UNETLoader", unet_name=unet_fixed_name, weight_dtype=unet_weight_dtype)
print(f"UNET node created with name: {unet_fixed_name}, weight_dtype: {unet_weight_dtype}")
# If we have a CLIP, load it with the appropriate CLIPLoader node.
if clip_info is not None:
@@ -538,23 +531,44 @@ class Sage_LoadModelwithGraphBuilderTest(ComfyNodeABC):
clip_path = clip_fixed_path if isinstance(clip_fixed_path, list) else [clip_fixed_path]
clip_type = clip_info["clip_type"]
print(f"Loading CLIP with type: {type_of_clip}, path: {clip_path}, clip_type: {clip_type}")
if type_of_clip == "Dual":
# Set clip_name1, clip_name2, and type
clip_node = graph.node("DualCLIPLoader", clip_name1=clip_path[0], clip_name2=clip_path[1], type=clip_type)
elif type_of_clip == "Triple":
# Set clip_name1, clip_name2, and clip_name3
clip_node = graph.node("TripleCLIPLoader", clip_name1=clip_path[0], clip_name2=clip_path[1], clip_name3=clip_path[2])
elif type_of_clip == "Quadruple":
# Set clip_name1, clip_name2, clip_name3, and clip_name4
clip_node = graph.node("QuadrupleCLIPLoader", clip_name1=clip_path[0], clip_name2=clip_path[1], clip_name3=clip_path[2], clip_name4=clip_path[3])
# If any of the clip paths are GGUF, we need to use the GGUF loader. It can also handle non-GGUF paths.
if any(get_file_extension(path) == "gguf" for path in clip_path):
try:
# If the CLIP is a GGUF, we need to use CLIPLoaderGGUF, DualCLIPLoaderGGUF, TripleCLIPLoaderGGUF, or QuadrupleCLIPLoaderGGUF nodes.
if type_of_clip == "Dual":
# Set clip_name1, clip_name2, and type
clip_node = graph.node("DualCLIPLoaderGGUF", clip_name1=clip_path[0], clip_name2=clip_path[1], type=clip_type)
elif type_of_clip == "Triple":
# Set clip_name1, clip_name2, and clip_name3
clip_node = graph.node("TripleCLIPLoaderGGUF", clip_name1=clip_path[0], clip_name2=clip_path[1], clip_name3=clip_path[2])
elif type_of_clip == "Quadruple":
# Set clip_name1, clip_name2, clip_name3, and clip_name4
clip_node = graph.node("QuadrupleCLIPLoaderGGUF", clip_name1=clip_path[0], clip_name2=clip_path[1], clip_name3=clip_path[2], clip_name4=clip_path[3])
else:
# Set clip_name and type
clip_node = graph.node("CLIPLoaderGGUF", clip_name=clip_path[0], type=clip_type)
except:
print("Unable to load CLIP as GGUF. Do you have ComfyUI-GGUF installed?")
raise ValueError("Unable to load CLIP as GGUF. Do you have ComfyUI-GGUF installed?")
else:
# Set clip_name and type
clip_node = graph.node("CLIPLoader", clip_name=clip_path[0], type=clip_type)
if type_of_clip == "Dual":
# Set clip_name1, clip_name2, and type
clip_node = graph.node("DualCLIPLoader", clip_name1=clip_path[0], clip_name2=clip_path[1], type=clip_type)
elif type_of_clip == "Triple":
# Set clip_name1, clip_name2, and clip_name3
clip_node = graph.node("TripleCLIPLoader", clip_name1=clip_path[0], clip_name2=clip_path[1], clip_name3=clip_path[2])
elif type_of_clip == "Quadruple":
# Set clip_name1, clip_name2, clip_name3, and clip_name4
clip_node = graph.node("QuadrupleCLIPLoader", clip_name1=clip_path[0], clip_name2=clip_path[1], clip_name3=clip_path[2], clip_name4=clip_path[3])
else:
# Set clip_name and type
clip_node = graph.node("CLIPLoader", clip_name=clip_path[0], type=clip_type)
# If we have a VAE, load it with the VAELoader node.
if vae_info is not None:
vae_name = vae_info["path"]
vae_fixed_name = get_path_without_base("vae", vae_name)
# There is no GGUF for VAE currently, so we can just use VAELoader.
vae_node = graph.node("VAELoader", vae_name=vae_fixed_name)
print(f"VAE node created with name: {vae_fixed_name}")
@@ -564,18 +578,23 @@ class Sage_LoadModelwithGraphBuilderTest(ComfyNodeABC):
unet_out = ckpt_node.out(0)
clip_out = ckpt_node.out(1)
vae_out = ckpt_node.out(2)
if unet_node is not None:
nodes.append(unet_node)
unet_out = unet_node.out(0)
if clip_node is not None:
nodes.append(clip_node)
clip_out = clip_node.out(0)
if vae_node is not None:
nodes.append(vae_node)
vae_out = vae_node.out(0)
print(f"Loaded model: {unet_out}, clip: {clip_out}, vae: {vae_out}")
return graph, unet_out, clip_out, vae_out
def load_model(self, model_info):
graph, unet_out, clip_out, vae_out = self.prepare_model_graph(model_info)
return {
"result": (unet_out, clip_out, vae_out),
"expand": graph.finalize()
+21 -6
View File
@@ -34,10 +34,13 @@ class Sage_CheckpointSelector(ComfyNodeABC):
class Sage_UNETSelector(ComfyNodeABC):
@classmethod
def INPUT_TYPES(cls) -> InputTypeDict:
unet_list = folder_paths.get_filename_list("diffusion_models")
unet_names = [x for x in folder_paths.get_filename_list("unet_gguf")]
unet_names += folder_paths.get_filename_list("diffusion_models")
unet_names = list(set(unet_names))
unet_names.sort() # Remove duplicates
return {
"required": {
"unet_name": (unet_list, {"tooltip": "The name of the UNET model to load."}),
"unet_name": (unet_names, {"tooltip": "The name of the UNET model to load."}),
"weight_dtype": (mi.weight_dtype_options, {"defaultInput": True, "default": "default", "tooltip": "The weight dtype to use for the UNET model."})
}
}
@@ -79,7 +82,10 @@ class Sage_VAESelector(ComfyNodeABC):
class Sage_CLIPSelector(ComfyNodeABC):
@classmethod
def INPUT_TYPES(cls) -> InputTypeDict:
model_list = folder_paths.get_filename_list("text_encoders")
model_list = [x for x in folder_paths.get_filename_list("text_encoders")]
model_list += folder_paths.get_filename_list("clip_gguf")
model_list = list(set(model_list))
model_list.sort() # Remove duplicates
return {
"required": {
"clip_name": (model_list, {"tooltip": "The name of the CLIP model to load."}),
@@ -102,7 +108,10 @@ class Sage_CLIPSelector(ComfyNodeABC):
class Sage_DualCLIPSelector(ComfyNodeABC):
@classmethod
def INPUT_TYPES(cls) -> InputTypeDict:
model_list = folder_paths.get_filename_list("text_encoders")
model_list = [x for x in folder_paths.get_filename_list("text_encoders")]
model_list += folder_paths.get_filename_list("clip_gguf")
model_list = list(set(model_list))
model_list.sort() # Remove duplicates
return {
"required": {
"clip_name_1": (model_list, {"tooltip": "The name of the first CLIP model to load."}),
@@ -126,7 +135,10 @@ class Sage_DualCLIPSelector(ComfyNodeABC):
class Sage_TripleCLIPSelector(ComfyNodeABC):
@classmethod
def INPUT_TYPES(cls) -> InputTypeDict:
model_list = folder_paths.get_filename_list("text_encoders")
model_list = [x for x in folder_paths.get_filename_list("text_encoders")]
model_list += folder_paths.get_filename_list("clip_gguf")
model_list = list(set(model_list))
model_list.sort() # Remove duplicates
return {
"required": {
"clip_name_1": (model_list, {"tooltip": "The name of the first CLIP model to load."}),
@@ -149,7 +161,10 @@ class Sage_TripleCLIPSelector(ComfyNodeABC):
class Sage_QuadCLIPSelector(ComfyNodeABC):
@classmethod
def INPUT_TYPES(cls) -> InputTypeDict:
model_list = folder_paths.get_filename_list("text_encoders")
model_list = [x for x in folder_paths.get_filename_list("text_encoders")]
model_list += folder_paths.get_filename_list("clip_gguf")
model_list = list(set(model_list))
model_list.sort() # Remove duplicates
return {
"required": {
"clip_name_1": (model_list, {"tooltip": "The name of the first CLIP model to load."}),