mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-09-10 18:57:06 +08:00
Allow default model paths to be re-used with other base paths.
- Update type hint to match actual usage. - Add re-usable function to add all default directories.
This commit is contained in:
parent
74325a5ecf
commit
fa3d6395e8
@ -4,7 +4,7 @@ import os
|
|||||||
import time
|
import time
|
||||||
import mimetypes
|
import mimetypes
|
||||||
import logging
|
import logging
|
||||||
from typing import Literal
|
from typing import Literal, TypedDict
|
||||||
from collections.abc import Collection
|
from collections.abc import Collection
|
||||||
|
|
||||||
supported_pt_extensions: set[str] = {'.ckpt', '.pt', '.bin', '.pth', '.safetensors', '.pkl', '.sft'}
|
supported_pt_extensions: set[str] = {'.ckpt', '.pt', '.bin', '.pth', '.safetensors', '.pkl', '.sft'}
|
||||||
@ -13,30 +13,75 @@ folder_names_and_paths: dict[str, tuple[list[str], set[str] | list[str]]] = {}
|
|||||||
|
|
||||||
base_path = os.path.dirname(os.path.realpath(__file__))
|
base_path = os.path.dirname(os.path.realpath(__file__))
|
||||||
models_dir = os.path.join(base_path, "models")
|
models_dir = os.path.join(base_path, "models")
|
||||||
folder_names_and_paths["checkpoints"] = ([os.path.join(models_dir, "checkpoints")], supported_pt_extensions)
|
|
||||||
folder_names_and_paths["configs"] = ([os.path.join(models_dir, "configs")], [".yaml"])
|
|
||||||
|
|
||||||
folder_names_and_paths["loras"] = ([os.path.join(models_dir, "loras")], supported_pt_extensions)
|
|
||||||
folder_names_and_paths["vae"] = ([os.path.join(models_dir, "vae")], supported_pt_extensions)
|
|
||||||
folder_names_and_paths["text_encoders"] = ([os.path.join(models_dir, "text_encoders"), os.path.join(models_dir, "clip")], supported_pt_extensions)
|
|
||||||
folder_names_and_paths["diffusion_models"] = ([os.path.join(models_dir, "unet"), os.path.join(models_dir, "diffusion_models")], supported_pt_extensions)
|
|
||||||
folder_names_and_paths["clip_vision"] = ([os.path.join(models_dir, "clip_vision")], supported_pt_extensions)
|
|
||||||
folder_names_and_paths["style_models"] = ([os.path.join(models_dir, "style_models")], supported_pt_extensions)
|
|
||||||
folder_names_and_paths["embeddings"] = ([os.path.join(models_dir, "embeddings")], supported_pt_extensions)
|
|
||||||
folder_names_and_paths["diffusers"] = ([os.path.join(models_dir, "diffusers")], ["folder"])
|
|
||||||
folder_names_and_paths["vae_approx"] = ([os.path.join(models_dir, "vae_approx")], supported_pt_extensions)
|
|
||||||
|
|
||||||
folder_names_and_paths["controlnet"] = ([os.path.join(models_dir, "controlnet"), os.path.join(models_dir, "t2i_adapter")], supported_pt_extensions)
|
class DefaultModelSubdirOptions(TypedDict, total=False):
|
||||||
folder_names_and_paths["gligen"] = ([os.path.join(models_dir, "gligen")], supported_pt_extensions)
|
extensions: set[str] | list[str]
|
||||||
|
alternate_names: list[str]
|
||||||
|
|
||||||
folder_names_and_paths["upscale_models"] = ([os.path.join(models_dir, "upscale_models")], supported_pt_extensions)
|
|
||||||
|
|
||||||
folder_names_and_paths["hypernetworks"] = ([os.path.join(models_dir, "hypernetworks")], supported_pt_extensions)
|
# Default models subdirs
|
||||||
|
default_paths: dict[str, DefaultModelSubdirOptions] = {
|
||||||
|
"checkpoints": {},
|
||||||
|
"configs": {
|
||||||
|
"extensions": [".yaml"],
|
||||||
|
},
|
||||||
|
"loras": {},
|
||||||
|
"vae": {},
|
||||||
|
"text_encoders": {
|
||||||
|
"alternate_names": ["clip"],
|
||||||
|
},
|
||||||
|
"clip_vision": {},
|
||||||
|
"style_models": {},
|
||||||
|
"embeddings": {},
|
||||||
|
"diffusers": {
|
||||||
|
"extensions": ["folder"],
|
||||||
|
},
|
||||||
|
"vae_approx": {},
|
||||||
|
"controlnet": {
|
||||||
|
"alternate_names": ["t2i_adapter"],
|
||||||
|
},
|
||||||
|
"gligen": {},
|
||||||
|
"upscale_models": {},
|
||||||
|
"hypernetworks": {},
|
||||||
|
"photomaker": {},
|
||||||
|
"classifiers": {
|
||||||
|
"extensions": {""},
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
folder_names_and_paths["photomaker"] = ([os.path.join(models_dir, "photomaker")], supported_pt_extensions)
|
def add_default_model_paths(models_path: str, is_default: bool = False) -> None:
|
||||||
|
"""
|
||||||
|
Adds all built-in model paths under a specified directory to the known folder paths.
|
||||||
|
|
||||||
folder_names_and_paths["classifiers"] = ([os.path.join(models_dir, "classifiers")], {""})
|
Args:
|
||||||
|
models_path: The path to the models directory that contains default folders, e.g. ``checkpoints``, ``loras``.
|
||||||
|
is_default: If the created paths should be the new default paths.
|
||||||
|
"""
|
||||||
|
|
||||||
|
for [name, options] in default_paths.items():
|
||||||
|
valid_extensions = options.get("extensions", supported_pt_extensions)
|
||||||
|
subdirs = [name, *options.get("alternate_names", [])]
|
||||||
|
|
||||||
|
for subdir in subdirs:
|
||||||
|
add_model_folder_path(
|
||||||
|
name,
|
||||||
|
os.path.join(models_path, subdir),
|
||||||
|
is_default=is_default,
|
||||||
|
is_legacy=True,
|
||||||
|
valid_extensions=valid_extensions,
|
||||||
|
)
|
||||||
|
|
||||||
|
# diffusion_models config shipped with the reverse order of other paths.
|
||||||
|
# Maintain original order, as this may impact ordering in extensions / UI.
|
||||||
|
add_model_folder_path("diffusion_models", os.path.join(models_path, "unet"), is_default=is_default, is_legacy=True, valid_extensions=supported_pt_extensions)
|
||||||
|
add_model_folder_path("diffusion_models", os.path.join(models_path, "diffusion_models"), is_default=is_default, is_legacy=True, valid_extensions=supported_pt_extensions)
|
||||||
|
|
||||||
|
|
||||||
|
# Default configuration: add model paths relative to this file.
|
||||||
|
add_default_model_paths(base_path)
|
||||||
|
|
||||||
|
# TODO: custom_nodes is now added out of original order; confirm no impact.
|
||||||
folder_names_and_paths["custom_nodes"] = ([os.path.join(base_path, "custom_nodes")], set())
|
folder_names_and_paths["custom_nodes"] = ([os.path.join(base_path, "custom_nodes")], set())
|
||||||
|
|
||||||
output_directory = os.path.join(base_path, "output")
|
output_directory = os.path.join(base_path, "output")
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user