feat: model_list node to display all the models available
This commit is contained in:
parent
ce92dd0570
commit
07926158f0
60
comfy-nodes/model_list.py
Normal file
60
comfy-nodes/model_list.py
Normal file
@ -0,0 +1,60 @@
|
||||
import folder_paths
|
||||
class AnyType(str):
|
||||
def __ne__(self, __value: object) -> bool:
|
||||
return False
|
||||
|
||||
from os import walk
|
||||
|
||||
WILDCARD = AnyType("*")
|
||||
|
||||
MODEL_EXTENSIONS = {
|
||||
"safetensors": "SafeTensors file format",
|
||||
"ckpt": "Checkpoint file",
|
||||
"pth": "PyTorch serialized file",
|
||||
"pkl": "Pickle file",
|
||||
"onnx": "ONNX file",
|
||||
}
|
||||
|
||||
def fetch_files(path):
|
||||
for (dirpath, dirnames, filenames) in walk(path):
|
||||
fs = []
|
||||
if len(dirnames) > 0:
|
||||
for dirname in dirnames:
|
||||
fs.extend(fetch_files(f"{dirpath}/{dirname}"))
|
||||
for filename in filenames:
|
||||
# Remove "./models/" from the beginning of dirpath
|
||||
relative_dirpath = dirpath.replace("./models/", "", 1)
|
||||
file_path = f"{relative_dirpath}/{filename}"
|
||||
|
||||
# Only add files that are known model extensions
|
||||
file_extension = filename.split('.')[-1].lower()
|
||||
if file_extension in MODEL_EXTENSIONS:
|
||||
fs.append(file_path)
|
||||
|
||||
return fs
|
||||
allModels = fetch_files("./models")
|
||||
|
||||
class ComfyUIDeployModalList:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"model": (allModels, ),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = (WILDCARD,)
|
||||
RETURN_NAMES = ("model",)
|
||||
|
||||
FUNCTION = "run"
|
||||
|
||||
CATEGORY = "model"
|
||||
|
||||
def run(self, model=""):
|
||||
# Split the model path by '/' and select the last item
|
||||
model_name = model.split('/')[-1]
|
||||
return [model_name]
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {"ComfyUIDeployModelList": ComfyUIDeployModalList}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {"ComfyUIDeployModelList": "Model List (ComfyUI Deploy)"}
|
Loading…
x
Reference in New Issue
Block a user