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