jinx/ComfyUI
1
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-08-02 23:14:49 +08:00
Code Issues Activity

Add folder_paths so models can be in multiple paths.

Browse Source
This commit is contained in:
comfyanonymous
2023-03-17 17:57:57 -04:00
parent 51bbbf8d64
commit e1a9e26968
3 changed files with 89 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
comfy_extras/nodes_upscale_model.py
View File

@@ -5,14 +5,12 @@ import model_management
from nodes import filter_files_extensions, recursive_search, supported_ckpt_extensions
import torch
import comfy.utils
import folder_paths
class UpscaleModelLoader:
models_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.realpath(__file__))), "models")
upscale_model_dir = os.path.join(models_dir, "upscale_models")
@classmethod
def INPUT_TYPES(s):
return {"required": { "model_name": (filter_files_extensions(recursive_search(s.upscale_model_dir), supported_ckpt_extensions), ),
return {"required": { "model_name": (folder_paths.get_filename_list("upscale_models"), ),
}}
RETURN_TYPES = ("UPSCALE_MODEL",)
FUNCTION = "load_model"
@@ -20,7 +18,7 @@ class UpscaleModelLoader:
CATEGORY = "loaders"
def load_model(self, model_name):
model_path = os.path.join(self.upscale_model_dir, model_name)
model_path = folder_paths.get_full_path("upscale_models", model_name)
sd = load_torch_file(model_path)
out = model_loading.load_state_dict(sd).eval()
return (out, )