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