mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-05-27 01:17:24 +08:00
Compare commits
14 Commits
4cd642517f
...
175659f7bd
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
175659f7bd | ||
|
|
f3ea976cba | ||
|
|
5538f62b0b | ||
|
|
2806163f6e | ||
|
|
cea8d0925f | ||
|
|
b138133ffa | ||
|
|
025e6792ee | ||
|
|
867b8d2408 | ||
|
|
d0f0b15cf5 | ||
|
|
b5bb83c964 | ||
|
|
f6d5068ac0 | ||
|
|
f334f2db3d | ||
|
|
e488f8bbc4 | ||
|
|
feee1c7a85 |
1
.gitignore
vendored
1
.gitignore
vendored
@ -8,6 +8,7 @@ __pycache__/
|
||||
/custom_nodes/
|
||||
!custom_nodes/example_node.py.example
|
||||
extra_model_paths.yaml
|
||||
extra_paths.yaml
|
||||
/.vs
|
||||
.vscode/
|
||||
.idea/
|
||||
|
||||
13
README.md
13
README.md
@ -1,7 +1,7 @@
|
||||
<div align="center">
|
||||
|
||||
# ComfyUI
|
||||
**The most powerful and modular visual AI engine and application.**
|
||||
**The most powerful and modular AI engine for content creation.**
|
||||
|
||||
|
||||
[![Website][website-shield]][website-url]
|
||||
@ -31,10 +31,16 @@
|
||||
[github-downloads-latest-shield]: https://img.shields.io/github/downloads/comfyanonymous/ComfyUI/latest/total?style=flat&label=downloads%40latest
|
||||
[github-downloads-link]: https://github.com/comfyanonymous/ComfyUI/releases
|
||||
|
||||

|
||||
<img width="1590" height="795" alt="ComfyUI Screenshot" src="https://github.com/user-attachments/assets/36e065e0-bfae-4456-8c7f-8369d5ea48a2" />
|
||||
<br>
|
||||
</div>
|
||||
|
||||
ComfyUI lets you design and execute advanced stable diffusion pipelines using a graph/nodes/flowchart based interface. Available on Windows, Linux, and macOS.
|
||||
ComfyUI is the AI creation engine for visual professionals who demand control over every model, every parameter, and every output. Its powerful and modular node graph interface empowers creatives to generate images, videos, 3D models, audio, and more...
|
||||
- ComfyUI natively supports the latest open-source state of the art models.
|
||||
- API nodes provide access to the best closed source models such as Nano Banana, Seedance, Hunyuan3D, etc.
|
||||
- It is available on Windows, Linux, and macOS, locally with our desktop application or on our cloud.
|
||||
- The most sophisticated workflows can be exposed through a simple UI thanks to App Mode.
|
||||
- It integrates seamlessly into production pipelines with our API endpoints.
|
||||
|
||||
## Get Started
|
||||
|
||||
@ -77,6 +83,7 @@ See what ComfyUI can do with the [newer template workflows](https://comfy.org/wo
|
||||
- [Hunyuan Image 2.1](https://comfyanonymous.github.io/ComfyUI_examples/hunyuan_image/)
|
||||
- [Flux 2](https://comfyanonymous.github.io/ComfyUI_examples/flux2/)
|
||||
- [Z Image](https://comfyanonymous.github.io/ComfyUI_examples/z_image/)
|
||||
- Ernie Image
|
||||
- Image Editing Models
|
||||
- [Omnigen 2](https://comfyanonymous.github.io/ComfyUI_examples/omnigen/)
|
||||
- [Flux Kontext](https://comfyanonymous.github.io/ComfyUI_examples/flux/#flux-kontext-image-editing-model)
|
||||
|
||||
@ -91,6 +91,7 @@ parser.add_argument("--directml", type=int, nargs="?", metavar="DIRECTML_DEVICE"
|
||||
|
||||
parser.add_argument("--oneapi-device-selector", type=str, default=None, metavar="SELECTOR_STRING", help="Sets the oneAPI device(s) this instance will use.")
|
||||
parser.add_argument("--supports-fp8-compute", action="store_true", help="ComfyUI will act like if the device supports fp8 compute.")
|
||||
parser.add_argument("--enable-triton-backend", action="store_true", help="ComfyUI will enable the use of Triton backend in comfy-kitchen. Is disabled at launch by default.")
|
||||
|
||||
class LatentPreviewMethod(enum.Enum):
|
||||
NoPreviews = "none"
|
||||
|
||||
@ -1,6 +1,8 @@
|
||||
import torch
|
||||
import logging
|
||||
|
||||
from comfy.cli_args import args
|
||||
|
||||
try:
|
||||
import comfy_kitchen as ck
|
||||
from comfy_kitchen.tensor import (
|
||||
@ -21,7 +23,15 @@ try:
|
||||
ck.registry.disable("cuda")
|
||||
logging.warning("WARNING: You need pytorch with cu130 or higher to use optimized CUDA operations.")
|
||||
|
||||
ck.registry.disable("triton")
|
||||
if args.enable_triton_backend:
|
||||
try:
|
||||
import triton
|
||||
logging.info("Found triton %s. Enabling comfy-kitchen triton backend.", triton.__version__)
|
||||
except ImportError as e:
|
||||
logging.error(f"Failed to import triton, Error: {e}, the comfy-kitchen triton backend will not be available.")
|
||||
ck.registry.disable("triton")
|
||||
else:
|
||||
ck.registry.disable("triton")
|
||||
for k, v in ck.list_backends().items():
|
||||
logging.info(f"Found comfy_kitchen backend {k}: {v}")
|
||||
except ImportError as e:
|
||||
|
||||
@ -202,14 +202,11 @@ class JoinImageWithAlpha(io.ComfyNode):
|
||||
|
||||
@classmethod
|
||||
def execute(cls, image: torch.Tensor, alpha: torch.Tensor) -> io.NodeOutput:
|
||||
batch_size = min(len(image), len(alpha))
|
||||
out_images = []
|
||||
|
||||
batch_size = max(len(image), len(alpha))
|
||||
alpha = 1.0 - resize_mask(alpha, image.shape[1:])
|
||||
for i in range(batch_size):
|
||||
out_images.append(torch.cat((image[i][:,:,:3], alpha[i].unsqueeze(2)), dim=2))
|
||||
|
||||
return io.NodeOutput(torch.stack(out_images))
|
||||
alpha = comfy.utils.repeat_to_batch_size(alpha, batch_size)
|
||||
image = comfy.utils.repeat_to_batch_size(image, batch_size)
|
||||
return io.NodeOutput(torch.cat((image[..., :3], alpha.unsqueeze(-1)), dim=-1))
|
||||
|
||||
|
||||
class CompositingExtension(ComfyExtension):
|
||||
|
||||
@ -666,12 +666,13 @@ class ColorTransfer(io.ComfyNode):
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="ColorTransfer",
|
||||
display_name="Color Transfer",
|
||||
category="image/postprocessing",
|
||||
description="Match the colors of one image to another using various algorithms.",
|
||||
search_aliases=["color match", "color grading", "color correction", "match colors", "color transform", "mkl", "reinhard", "histogram"],
|
||||
inputs=[
|
||||
io.Image.Input("image_target", tooltip="Image(s) to apply the color transform to."),
|
||||
io.Image.Input("image_ref", optional=True, tooltip="Reference image(s) to match colors to. If not provided, processing is skipped"),
|
||||
io.Image.Input("image_ref", tooltip="Reference image(s) to match colors to."),
|
||||
io.Combo.Input("method", options=['reinhard_lab', 'mkl_lab', 'histogram'],),
|
||||
io.DynamicCombo.Input("source_stats",
|
||||
tooltip="per_frame: each frame matched to image_ref individually. uniform: pool stats across all source frames as baseline, match to image_ref. target_frame: use one chosen frame as the baseline for the transform to image_ref, applied uniformly to all frames (preserves relative differences)",
|
||||
|
||||
@ -49,7 +49,7 @@ class Int(io.ComfyNode):
|
||||
display_name="Int",
|
||||
category="utils/primitive",
|
||||
inputs=[
|
||||
io.Int.Input("value", min=-sys.maxsize, max=sys.maxsize, control_after_generate=True),
|
||||
io.Int.Input("value", min=-sys.maxsize, max=sys.maxsize, control_after_generate=io.ControlAfterGenerate.fixed),
|
||||
],
|
||||
outputs=[io.Int.Output()],
|
||||
)
|
||||
|
||||
@ -1,8 +1,16 @@
|
||||
#Rename this to extra_model_paths.yaml and ComfyUI will load it
|
||||
#
|
||||
# DEPRECATED: extra_model_paths.yaml is superseded by extra_paths.yaml, which supports
|
||||
# all path configuration (system dirs, custom_nodes, and models) in a cleaner format.
|
||||
# See extra_paths.yaml.example. This file continues to work for backward compatibility.
|
||||
|
||||
#config for comfyui
|
||||
#your base path should be either an existing comfy install or a central folder where you store all of your models, loras, etc.
|
||||
|
||||
# When base_path is set, any standard subdirectory that exists on disk is automatically
|
||||
# registered — the explicit paths below are optional and only needed to override a
|
||||
# category or point it to a non-standard location.
|
||||
|
||||
#comfyui:
|
||||
# base_path: path/to/comfyui/
|
||||
# # You can use is_default to mark that these folders should be listed first, and used as the default dirs for eg downloads
|
||||
@ -28,7 +36,7 @@
|
||||
#config for a1111 ui
|
||||
#all you have to do is uncomment this (remove the #) and change the base_path to where yours is installed
|
||||
|
||||
#a111:
|
||||
#a1111:
|
||||
# base_path: path/to/stable-diffusion-webui/
|
||||
# checkpoints: models/Stable-diffusion
|
||||
# configs: models/Stable-diffusion
|
||||
|
||||
84
extra_paths.yaml.example
Normal file
84
extra_paths.yaml.example
Normal file
@ -0,0 +1,84 @@
|
||||
#Rename this to extra_paths.yaml and ComfyUI will load it
|
||||
#This is the successor to extra_model_paths.yaml and supports all path configuration in one file.
|
||||
|
||||
#config for comfyui
|
||||
#Set base_path to your ComfyUI install root. System directories (output, input, temp, user)
|
||||
#and custom_nodes are resolved relative to base_path.
|
||||
#
|
||||
#Model paths go under the 'models' block. If you only set models/base_path, all standard
|
||||
#subdirectories that exist on disk are automatically registered — no need to list them.
|
||||
#Explicit paths under models/ are optional and only needed to override a specific category
|
||||
#or point it to a non-standard location.
|
||||
|
||||
#comfyui:
|
||||
# base_path: path/to/comfyui/
|
||||
# # System directories (relative to base_path, or absolute)
|
||||
# output: output/
|
||||
# input: input/
|
||||
# temp: temp/
|
||||
# user: user/
|
||||
# # Custom nodes directory (not auto-scanned; explicit only)
|
||||
# custom_nodes: custom_nodes/
|
||||
# models:
|
||||
# base_path: models/
|
||||
# # You can use is_default to mark that these folders should be listed first,
|
||||
# # and used as the default dirs for eg downloads
|
||||
# #is_default: true
|
||||
# checkpoints: checkpoints/
|
||||
# text_encoders: |
|
||||
# text_encoders/
|
||||
# clip/ # legacy location still supported
|
||||
# clip_vision: clip_vision/
|
||||
# configs: configs/
|
||||
# controlnet: |
|
||||
# controlnet/
|
||||
# t2i_adapter/
|
||||
# diffusion_models: |
|
||||
# diffusion_models/
|
||||
# unet/
|
||||
# diffusers: diffusers/
|
||||
# embeddings: embeddings/
|
||||
# frame_interpolation: frame_interpolation/
|
||||
# gligen: gligen/
|
||||
# hypernetworks: hypernetworks/
|
||||
# latent_upscale_models: latent_upscale_models/
|
||||
# loras: loras/
|
||||
# model_patches: model_patches/
|
||||
# photomaker: photomaker/
|
||||
# style_models: style_models/
|
||||
# upscale_models: upscale_models/
|
||||
# vae: vae/
|
||||
# vae_approx: vae_approx/
|
||||
# audio_encoders: audio_encoders/
|
||||
# classifiers: classifiers/
|
||||
|
||||
|
||||
#config for a1111 ui
|
||||
#all you have to do is uncomment this (remove the #) and change the base_path to where yours is installed
|
||||
|
||||
#a111:
|
||||
# models:
|
||||
# base_path: path/to/stable-diffusion-webui/
|
||||
# checkpoints: models/Stable-diffusion
|
||||
# configs: models/Stable-diffusion
|
||||
# vae: models/VAE
|
||||
# loras: |
|
||||
# models/Lora
|
||||
# models/LyCORIS
|
||||
# upscale_models: |
|
||||
# models/ESRGAN
|
||||
# models/RealESRGAN
|
||||
# models/SwinIR
|
||||
# embeddings: embeddings
|
||||
# hypernetworks: models/hypernetworks
|
||||
# controlnet: models/ControlNet
|
||||
|
||||
|
||||
# For a full list of supported model category keys (style_models, vae_approx, hypernetworks,
|
||||
# photomaker, model_patches, audio_encoders, classifiers, etc.) see folder_paths.py.
|
||||
|
||||
#other_ui:
|
||||
# models:
|
||||
# base_path: path/to/ui
|
||||
# checkpoints: models/checkpoints
|
||||
# gligen: models/gligen
|
||||
18
main.py
18
main.py
@ -93,9 +93,21 @@ if args.enable_manager:
|
||||
|
||||
|
||||
def apply_custom_paths():
|
||||
# extra model paths
|
||||
extra_model_paths_config_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "extra_model_paths.yaml")
|
||||
if os.path.isfile(extra_model_paths_config_path):
|
||||
install_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
|
||||
# extra_paths.yaml — primary config (superset of extra_model_paths.yaml)
|
||||
extra_paths_config_path = os.path.join(install_dir, "extra_paths.yaml")
|
||||
extra_model_paths_config_path = os.path.join(install_dir, "extra_model_paths.yaml")
|
||||
|
||||
if os.path.isfile(extra_paths_config_path):
|
||||
utils.extra_config.load_extra_path_config(extra_paths_config_path, allow_system_dirs=True)
|
||||
if os.path.isfile(extra_model_paths_config_path):
|
||||
logging.warning(
|
||||
"Both extra_paths.yaml and extra_model_paths.yaml found; "
|
||||
"ignoring the deprecated extra_model_paths.yaml. "
|
||||
"Please remove or migrate its entries to extra_paths.yaml."
|
||||
)
|
||||
elif os.path.isfile(extra_model_paths_config_path):
|
||||
utils.extra_config.load_extra_path_config(extra_model_paths_config_path)
|
||||
|
||||
if args.extra_model_paths_config:
|
||||
|
||||
@ -86,6 +86,6 @@ def image_alpha_fix(destination, source):
|
||||
if destination.shape[-1] < source.shape[-1]:
|
||||
source = source[...,:destination.shape[-1]]
|
||||
elif destination.shape[-1] > source.shape[-1]:
|
||||
destination = torch.nn.functional.pad(destination, (0, 1))
|
||||
destination[..., -1] = 1.0
|
||||
source = torch.nn.functional.pad(source, (0, 1))
|
||||
source[..., -1] = 1.0
|
||||
return destination, source
|
||||
|
||||
66
nodes.py
66
nodes.py
@ -1754,57 +1754,49 @@ class LoadImage:
|
||||
|
||||
return True
|
||||
|
||||
class LoadImageMask:
|
||||
|
||||
class LoadImageMask(LoadImage):
|
||||
ESSENTIALS_CATEGORY = "Image Tools"
|
||||
SEARCH_ALIASES = ["import mask", "alpha mask", "channel mask"]
|
||||
|
||||
_color_channels = ["alpha", "red", "green", "blue"]
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
input_dir = folder_paths.get_input_directory()
|
||||
files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
|
||||
return {"required":
|
||||
{"image": (sorted(files), {"image_upload": True}),
|
||||
"channel": (s._color_channels, ), }
|
||||
}
|
||||
types = super().INPUT_TYPES()
|
||||
return {
|
||||
"required": {
|
||||
**types["required"],
|
||||
"channel": (s._color_channels, )
|
||||
}
|
||||
}
|
||||
|
||||
CATEGORY = "mask"
|
||||
|
||||
RETURN_TYPES = ("MASK",)
|
||||
FUNCTION = "load_image"
|
||||
def load_image(self, image, channel):
|
||||
image_path = folder_paths.get_annotated_filepath(image)
|
||||
i = node_helpers.pillow(Image.open, image_path)
|
||||
i = node_helpers.pillow(ImageOps.exif_transpose, i)
|
||||
if i.getbands() != ("R", "G", "B", "A"):
|
||||
if i.mode == 'I':
|
||||
i = i.point(lambda i: i * (1 / 255))
|
||||
i = i.convert("RGBA")
|
||||
mask = None
|
||||
FUNCTION = "load_image_mask"
|
||||
|
||||
def load_image_mask(self, image, channel):
|
||||
image_tensor, mask_tensor = super().load_image(image)
|
||||
c = channel[0].upper()
|
||||
if c in i.getbands():
|
||||
mask = np.array(i.getchannel(c)).astype(np.float32) / 255.0
|
||||
mask = torch.from_numpy(mask)
|
||||
if c == 'A':
|
||||
mask = 1. - mask
|
||||
|
||||
if c == 'A':
|
||||
return (mask_tensor,)
|
||||
|
||||
channel_idx = {'R': 0, 'G': 1, 'B': 2}.get(c, 0)
|
||||
|
||||
if channel_idx < image_tensor.shape[-1]:
|
||||
return (image_tensor[..., channel_idx].clone(),)
|
||||
else:
|
||||
mask = torch.zeros((64,64), dtype=torch.float32, device="cpu")
|
||||
return (mask.unsqueeze(0),)
|
||||
empty_mask = torch.zeros(
|
||||
image_tensor.shape[:-1],
|
||||
dtype=image_tensor.dtype,
|
||||
device=image_tensor.device
|
||||
)
|
||||
return (empty_mask,)
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(s, image, channel):
|
||||
image_path = folder_paths.get_annotated_filepath(image)
|
||||
m = hashlib.sha256()
|
||||
with open(image_path, 'rb') as f:
|
||||
m.update(f.read())
|
||||
return m.digest().hex()
|
||||
|
||||
@classmethod
|
||||
def VALIDATE_INPUTS(s, image):
|
||||
if not folder_paths.exists_annotated_filepath(image):
|
||||
return "Invalid image file: {}".format(image)
|
||||
|
||||
return True
|
||||
return super().IS_CHANGED(image)
|
||||
|
||||
|
||||
class LoadImageOutput(LoadImage):
|
||||
|
||||
@ -1,3 +1,4 @@
|
||||
import errno
|
||||
import os
|
||||
import sys
|
||||
import asyncio
|
||||
@ -1245,7 +1246,13 @@ class PromptServer():
|
||||
address = addr[0]
|
||||
port = addr[1]
|
||||
site = web.TCPSite(runner, address, port, ssl_context=ssl_ctx)
|
||||
await site.start()
|
||||
try:
|
||||
await site.start()
|
||||
except OSError as e:
|
||||
if e.errno == errno.EADDRINUSE:
|
||||
logging.error(f"Port {port} is already in use on address {address}. Please close the other application or use a different port with --port.")
|
||||
raise SystemExit(1)
|
||||
raise
|
||||
|
||||
if not hasattr(self, 'address'):
|
||||
self.address = address #TODO: remove this
|
||||
|
||||
@ -301,3 +301,320 @@ def test_load_extra_path_config_no_base_path(
|
||||
actual_diffusion = folder_paths.folder_names_and_paths["diffusion_models"][0]
|
||||
assert len(actual_diffusion) == 1, "Should have one path for 'diffusion_models'."
|
||||
assert actual_diffusion[0] == os.path.abspath(expected_unet)
|
||||
|
||||
|
||||
@patch("yaml.safe_load")
|
||||
def test_load_extra_path_config_implicit_subdirs(
|
||||
mock_yaml_load, clear_folder_paths, tmp_path
|
||||
):
|
||||
"""
|
||||
When base_path is set and no explicit sub-paths are declared, any subdir
|
||||
whose name matches a known category is auto-registered.
|
||||
"""
|
||||
# Create real subdirs that match known categories
|
||||
(tmp_path / "checkpoints").mkdir()
|
||||
(tmp_path / "loras").mkdir()
|
||||
(tmp_path / "unknown_dir").mkdir() # not a registered category — should be ignored
|
||||
|
||||
config_data = {
|
||||
"comfyui": {
|
||||
"base_path": str(tmp_path),
|
||||
}
|
||||
}
|
||||
mock_yaml_load.return_value = config_data
|
||||
|
||||
# Pre-populate only the categories we're testing so clear_folder_paths doesn't hide them
|
||||
folder_paths.folder_names_and_paths["checkpoints"] = ([], set())
|
||||
folder_paths.folder_names_and_paths["loras"] = ([], set())
|
||||
|
||||
yaml_path = str(tmp_path / "extra_model_paths.yaml")
|
||||
with open(yaml_path, "w") as f:
|
||||
f.write("") # content ignored; yaml.safe_load is mocked
|
||||
|
||||
load_extra_path_config(yaml_path)
|
||||
|
||||
assert str(tmp_path / "checkpoints") in folder_paths.folder_names_and_paths["checkpoints"][0]
|
||||
assert str(tmp_path / "loras") in folder_paths.folder_names_and_paths["loras"][0]
|
||||
assert "unknown_dir" not in folder_paths.folder_names_and_paths
|
||||
|
||||
|
||||
@patch("yaml.safe_load")
|
||||
def test_implicit_scan_excludes_custom_nodes(
|
||||
mock_yaml_load, clear_folder_paths, tmp_path
|
||||
):
|
||||
"""custom_nodes must never be auto-registered by the implicit scan."""
|
||||
(tmp_path / "custom_nodes").mkdir()
|
||||
(tmp_path / "checkpoints").mkdir()
|
||||
|
||||
config_data = {"comfyui": {"base_path": str(tmp_path)}}
|
||||
mock_yaml_load.return_value = config_data
|
||||
|
||||
folder_paths.folder_names_and_paths["checkpoints"] = ([], set())
|
||||
folder_paths.folder_names_and_paths["custom_nodes"] = ([], set())
|
||||
|
||||
yaml_path = str(tmp_path / "extra_paths.yaml")
|
||||
with open(yaml_path, "w") as f:
|
||||
f.write("")
|
||||
|
||||
load_extra_path_config(yaml_path)
|
||||
|
||||
assert str(tmp_path / "checkpoints") in folder_paths.folder_names_and_paths["checkpoints"][0]
|
||||
assert str(tmp_path / "custom_nodes") not in folder_paths.folder_names_and_paths["custom_nodes"][0], \
|
||||
"custom_nodes must not be auto-registered by the implicit scan"
|
||||
|
||||
|
||||
@patch("yaml.safe_load")
|
||||
def test_load_extra_path_config_explicit_overrides_implicit(
|
||||
mock_yaml_load, clear_folder_paths, tmp_path
|
||||
):
|
||||
"""
|
||||
Explicit sub-path declarations take precedence; the implicit scan must not
|
||||
double-register a category that was already declared explicitly.
|
||||
"""
|
||||
(tmp_path / "loras").mkdir()
|
||||
custom_loras = tmp_path / "my_custom_loras"
|
||||
custom_loras.mkdir()
|
||||
|
||||
config_data = {
|
||||
"comfyui": {
|
||||
"base_path": str(tmp_path),
|
||||
"loras": "my_custom_loras", # explicit override
|
||||
}
|
||||
}
|
||||
mock_yaml_load.return_value = config_data
|
||||
folder_paths.folder_names_and_paths["loras"] = ([], set())
|
||||
|
||||
yaml_path = str(tmp_path / "extra_model_paths.yaml")
|
||||
with open(yaml_path, "w") as f:
|
||||
f.write("")
|
||||
|
||||
load_extra_path_config(yaml_path)
|
||||
|
||||
registered = folder_paths.folder_names_and_paths["loras"][0]
|
||||
assert str(custom_loras) in registered
|
||||
assert str(tmp_path / "loras") not in registered, "Implicit path must not override explicit"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def save_restore_system_dirs():
|
||||
"""Save and restore folder_paths system directories around a test."""
|
||||
saved = {
|
||||
"output": folder_paths.get_output_directory(),
|
||||
"input": folder_paths.get_input_directory(),
|
||||
"temp": folder_paths.get_temp_directory(),
|
||||
"user": folder_paths.get_user_directory(),
|
||||
}
|
||||
yield
|
||||
folder_paths.set_output_directory(saved["output"])
|
||||
folder_paths.set_input_directory(saved["input"])
|
||||
folder_paths.set_temp_directory(saved["temp"])
|
||||
folder_paths.set_user_directory(saved["user"])
|
||||
|
||||
|
||||
@patch("yaml.safe_load")
|
||||
def test_system_dir_keys(mock_yaml_load, save_restore_system_dirs, tmp_path):
|
||||
"""System directory keys (output, input, temp, user) call set_*_directory()."""
|
||||
config_data = {
|
||||
"comfyui": {
|
||||
"base_path": str(tmp_path),
|
||||
"output": "my_output/",
|
||||
"input": "my_input/",
|
||||
"temp": "my_temp/",
|
||||
"user": "my_user/",
|
||||
}
|
||||
}
|
||||
mock_yaml_load.return_value = config_data
|
||||
|
||||
yaml_path = str(tmp_path / "extra_paths.yaml")
|
||||
with open(yaml_path, "w") as f:
|
||||
f.write("")
|
||||
|
||||
load_extra_path_config(yaml_path, allow_system_dirs=True)
|
||||
|
||||
assert folder_paths.get_output_directory() == os.path.normpath(str(tmp_path / "my_output"))
|
||||
assert folder_paths.get_input_directory() == os.path.normpath(str(tmp_path / "my_input"))
|
||||
assert folder_paths.get_temp_directory() == os.path.normpath(str(tmp_path / "my_temp"))
|
||||
assert folder_paths.get_user_directory() == os.path.normpath(str(tmp_path / "my_user"))
|
||||
|
||||
|
||||
@patch("yaml.safe_load")
|
||||
def test_system_dir_keys_not_applied_for_legacy(mock_yaml_load, save_restore_system_dirs, tmp_path):
|
||||
"""System directory keys are ignored when allow_system_dirs=False (legacy extra_model_paths.yaml)."""
|
||||
original_output = folder_paths.get_output_directory()
|
||||
config_data = {
|
||||
"comfyui": {
|
||||
"base_path": str(tmp_path),
|
||||
"output": "my_output/",
|
||||
}
|
||||
}
|
||||
mock_yaml_load.return_value = config_data
|
||||
|
||||
yaml_path = str(tmp_path / "extra_model_paths.yaml")
|
||||
with open(yaml_path, "w") as f:
|
||||
f.write("")
|
||||
|
||||
load_extra_path_config(yaml_path) # allow_system_dirs defaults to False
|
||||
|
||||
# output should be unchanged — treated as a model category, not a system dir
|
||||
assert folder_paths.get_output_directory() == original_output
|
||||
|
||||
|
||||
@patch("yaml.safe_load")
|
||||
def test_nested_models_block(mock_yaml_load, clear_folder_paths, tmp_path):
|
||||
"""Nested models: block registers model paths relative to models/base_path."""
|
||||
config_data = {
|
||||
"comfyui": {
|
||||
"base_path": str(tmp_path),
|
||||
"models": {
|
||||
"base_path": "models/",
|
||||
"checkpoints": "checkpoints/",
|
||||
"loras": "loras/",
|
||||
},
|
||||
}
|
||||
}
|
||||
mock_yaml_load.return_value = config_data
|
||||
|
||||
folder_paths.folder_names_and_paths["checkpoints"] = ([], set())
|
||||
folder_paths.folder_names_and_paths["loras"] = ([], set())
|
||||
|
||||
yaml_path = str(tmp_path / "extra_paths.yaml")
|
||||
with open(yaml_path, "w") as f:
|
||||
f.write("")
|
||||
|
||||
load_extra_path_config(yaml_path)
|
||||
|
||||
expected_ckpt = os.path.normpath(str(tmp_path / "models" / "checkpoints"))
|
||||
expected_loras = os.path.normpath(str(tmp_path / "models" / "loras"))
|
||||
assert expected_ckpt in folder_paths.folder_names_and_paths["checkpoints"][0]
|
||||
assert expected_loras in folder_paths.folder_names_and_paths["loras"][0]
|
||||
|
||||
|
||||
@patch("yaml.safe_load")
|
||||
def test_nested_models_is_default(mock_yaml_load, clear_folder_paths, tmp_path):
|
||||
"""is_default under models: applies to all model paths in that block."""
|
||||
config_data = {
|
||||
"comfyui": {
|
||||
"models": {
|
||||
"base_path": str(tmp_path),
|
||||
"is_default": True,
|
||||
"checkpoints": "checkpoints/",
|
||||
},
|
||||
}
|
||||
}
|
||||
mock_yaml_load.return_value = config_data
|
||||
folder_paths.folder_names_and_paths["checkpoints"] = ([], set())
|
||||
|
||||
yaml_path = str(tmp_path / "extra_paths.yaml")
|
||||
with open(yaml_path, "w") as f:
|
||||
f.write("")
|
||||
|
||||
mock_add = Mock()
|
||||
with patch.object(folder_paths, "add_model_folder_path", mock_add):
|
||||
load_extra_path_config(yaml_path)
|
||||
|
||||
call = mock_add.call_args_list[0]
|
||||
assert call.args[0] == "checkpoints"
|
||||
assert call.args[2] is True, "is_default under models: must be passed as True"
|
||||
|
||||
|
||||
@patch("yaml.safe_load")
|
||||
def test_nested_models_multipath(mock_yaml_load, clear_folder_paths, tmp_path):
|
||||
"""Multi-line path values inside models: register multiple paths per category."""
|
||||
config_data = {
|
||||
"comfyui": {
|
||||
"models": {
|
||||
"base_path": str(tmp_path),
|
||||
"text_encoders": "text_encoders/\nclip/",
|
||||
},
|
||||
}
|
||||
}
|
||||
mock_yaml_load.return_value = config_data
|
||||
folder_paths.folder_names_and_paths["text_encoders"] = ([], set())
|
||||
|
||||
yaml_path = str(tmp_path / "extra_paths.yaml")
|
||||
with open(yaml_path, "w") as f:
|
||||
f.write("")
|
||||
|
||||
load_extra_path_config(yaml_path)
|
||||
|
||||
registered = folder_paths.folder_names_and_paths["text_encoders"][0]
|
||||
assert os.path.normpath(str(tmp_path / "text_encoders")) in registered
|
||||
assert os.path.normpath(str(tmp_path / "clip")) in registered
|
||||
|
||||
|
||||
@patch("yaml.safe_load")
|
||||
def test_nested_models_auto_scan(mock_yaml_load, clear_folder_paths, tmp_path):
|
||||
"""models: with only base_path auto-scans for known categories that exist on disk."""
|
||||
(tmp_path / "models" / "checkpoints").mkdir(parents=True)
|
||||
(tmp_path / "models" / "loras").mkdir()
|
||||
|
||||
config_data = {
|
||||
"comfyui": {
|
||||
"base_path": str(tmp_path),
|
||||
"models": {"base_path": "models/"},
|
||||
}
|
||||
}
|
||||
mock_yaml_load.return_value = config_data
|
||||
folder_paths.folder_names_and_paths["checkpoints"] = ([], set())
|
||||
folder_paths.folder_names_and_paths["loras"] = ([], set())
|
||||
|
||||
yaml_path = str(tmp_path / "extra_paths.yaml")
|
||||
with open(yaml_path, "w") as f:
|
||||
f.write("")
|
||||
|
||||
load_extra_path_config(yaml_path)
|
||||
|
||||
assert os.path.normpath(str(tmp_path / "models" / "checkpoints")) in \
|
||||
folder_paths.folder_names_and_paths["checkpoints"][0]
|
||||
assert os.path.normpath(str(tmp_path / "models" / "loras")) in \
|
||||
folder_paths.folder_names_and_paths["loras"][0]
|
||||
|
||||
|
||||
@patch("yaml.safe_load")
|
||||
def test_explicit_custom_nodes_key(mock_yaml_load, clear_folder_paths, tmp_path):
|
||||
"""Explicit custom_nodes key in a block registers the path via add_model_folder_path."""
|
||||
config_data = {
|
||||
"comfyui": {
|
||||
"base_path": str(tmp_path),
|
||||
"custom_nodes": "my_nodes/",
|
||||
}
|
||||
}
|
||||
mock_yaml_load.return_value = config_data
|
||||
folder_paths.folder_names_and_paths["custom_nodes"] = ([], set())
|
||||
|
||||
yaml_path = str(tmp_path / "extra_paths.yaml")
|
||||
with open(yaml_path, "w") as f:
|
||||
f.write("")
|
||||
|
||||
load_extra_path_config(yaml_path)
|
||||
|
||||
assert os.path.normpath(str(tmp_path / "my_nodes")) in \
|
||||
folder_paths.folder_names_and_paths["custom_nodes"][0]
|
||||
|
||||
|
||||
@patch("yaml.safe_load")
|
||||
def test_nested_models_inherits_block_base(mock_yaml_load, clear_folder_paths, tmp_path):
|
||||
"""models: block without its own base_path inherits the outer block's base_path."""
|
||||
config_data = {
|
||||
"comfyui": {
|
||||
"base_path": str(tmp_path),
|
||||
"is_default": True,
|
||||
"models": {
|
||||
"checkpoints": "models/checkpoints/",
|
||||
},
|
||||
}
|
||||
}
|
||||
mock_yaml_load.return_value = config_data
|
||||
folder_paths.folder_names_and_paths["checkpoints"] = ([], set())
|
||||
|
||||
yaml_path = str(tmp_path / "extra_paths.yaml")
|
||||
with open(yaml_path, "w") as f:
|
||||
f.write("")
|
||||
|
||||
load_extra_path_config(yaml_path)
|
||||
|
||||
expected = os.path.normpath(str(tmp_path / "models" / "checkpoints"))
|
||||
paths = folder_paths.folder_names_and_paths["checkpoints"][0]
|
||||
assert expected in paths
|
||||
# is_default inherited: path should be at index 0
|
||||
assert paths[0] == expected
|
||||
|
||||
@ -1,34 +1,101 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import yaml
|
||||
import folder_paths
|
||||
import logging
|
||||
|
||||
def load_extra_path_config(yaml_path):
|
||||
_SYSTEM_DIR_KEYS = frozenset({"output", "input", "temp", "user"})
|
||||
|
||||
|
||||
def _resolve_base(raw: str, parent_base: str | None, yaml_dir: str) -> str:
|
||||
"""Resolve a base_path value: expand vars/user, join onto parent_base or yaml_dir if relative."""
|
||||
raw = os.path.expandvars(os.path.expanduser(raw))
|
||||
if not os.path.isabs(raw):
|
||||
anchor = parent_base if parent_base else yaml_dir
|
||||
raw = os.path.abspath(os.path.join(anchor, raw))
|
||||
return os.path.normpath(raw)
|
||||
|
||||
|
||||
def _add_model_paths(category: str, raw_value: str, base: str | None, yaml_dir: str, is_default: bool) -> None:
|
||||
"""Split a (possibly multi-line) path value and register each path as a model folder."""
|
||||
for raw in str(raw_value).split("\n"):
|
||||
raw = raw.strip()
|
||||
if not raw:
|
||||
continue
|
||||
if base and not os.path.isabs(raw):
|
||||
full_path = os.path.join(base, raw)
|
||||
elif not os.path.isabs(raw):
|
||||
full_path = os.path.abspath(os.path.join(yaml_dir, raw))
|
||||
else:
|
||||
full_path = raw
|
||||
normalized = os.path.normpath(full_path)
|
||||
logging.info("Adding extra search path %s %s", category, normalized)
|
||||
folder_paths.add_model_folder_path(category, normalized, is_default)
|
||||
|
||||
|
||||
def _implicit_scan(base: str, exclude: set[str], is_default: bool) -> None:
|
||||
"""Auto-register base/<category>/ for known model categories that exist on disk.
|
||||
|
||||
custom_nodes and system directory keys are always excluded from the scan.
|
||||
"""
|
||||
skip = _SYSTEM_DIR_KEYS | {"custom_nodes"} | exclude
|
||||
for category in folder_paths.folder_names_and_paths:
|
||||
if category in skip:
|
||||
continue
|
||||
path = os.path.normpath(os.path.join(base, category))
|
||||
if os.path.isdir(path):
|
||||
logging.info("Adding extra search path %s %s", category, path)
|
||||
folder_paths.add_model_folder_path(category, path, is_default)
|
||||
|
||||
|
||||
def load_extra_path_config(yaml_path: str, allow_system_dirs: bool = False) -> None:
|
||||
with open(yaml_path, 'r', encoding='utf-8') as stream:
|
||||
config = yaml.safe_load(stream)
|
||||
yaml_dir = os.path.dirname(os.path.abspath(yaml_path))
|
||||
for c in config:
|
||||
conf = config[c]
|
||||
|
||||
for _block_name, conf in config.items():
|
||||
if conf is None:
|
||||
continue
|
||||
base_path = None
|
||||
|
||||
# Pop block-level meta keys (preserved for flat backward-compat style)
|
||||
block_base = None
|
||||
if "base_path" in conf:
|
||||
base_path = conf.pop("base_path")
|
||||
base_path = os.path.expandvars(os.path.expanduser(base_path))
|
||||
if not os.path.isabs(base_path):
|
||||
base_path = os.path.abspath(os.path.join(yaml_dir, base_path))
|
||||
is_default = False
|
||||
if "is_default" in conf:
|
||||
is_default = conf.pop("is_default")
|
||||
for x in conf:
|
||||
for y in conf[x].split("\n"):
|
||||
if len(y) == 0:
|
||||
continue
|
||||
full_path = y
|
||||
if base_path:
|
||||
full_path = os.path.join(base_path, full_path)
|
||||
elif not os.path.isabs(full_path):
|
||||
full_path = os.path.abspath(os.path.join(yaml_dir, y))
|
||||
normalized_path = os.path.normpath(full_path)
|
||||
logging.info("Adding extra search path {} {}".format(x, normalized_path))
|
||||
folder_paths.add_model_folder_path(x, normalized_path, is_default)
|
||||
block_base = _resolve_base(conf.pop("base_path"), None, yaml_dir)
|
||||
block_is_default = bool(conf.pop("is_default", False))
|
||||
|
||||
has_models_block = False
|
||||
flat_model_keys: set[str] = set()
|
||||
|
||||
for key, value in conf.items():
|
||||
if allow_system_dirs and key in _SYSTEM_DIR_KEYS:
|
||||
# System directory override → set_*_directory()
|
||||
path = _resolve_base(str(value).strip(), block_base, yaml_dir)
|
||||
logging.info("Setting %s directory to %s", key, path)
|
||||
getattr(folder_paths, f"set_{key}_directory")(path)
|
||||
|
||||
elif key == "custom_nodes":
|
||||
_add_model_paths("custom_nodes", value, block_base, yaml_dir, block_is_default)
|
||||
|
||||
elif key == "models" and isinstance(value, dict):
|
||||
# New nested style: models: { base_path, is_default, <categories> }
|
||||
has_models_block = True
|
||||
models_conf = dict(value)
|
||||
models_base = block_base
|
||||
if "base_path" in models_conf:
|
||||
models_base = _resolve_base(models_conf.pop("base_path"), block_base, yaml_dir)
|
||||
models_is_default = bool(models_conf.pop("is_default", block_is_default))
|
||||
explicit: set[str] = set(models_conf.keys())
|
||||
for cat, raw in models_conf.items():
|
||||
_add_model_paths(cat, raw, models_base, yaml_dir, models_is_default)
|
||||
if models_base:
|
||||
_implicit_scan(models_base, explicit, models_is_default)
|
||||
|
||||
else:
|
||||
# Flat model key — backward-compat style
|
||||
_add_model_paths(key, value, block_base, yaml_dir, block_is_default)
|
||||
flat_model_keys.add(key)
|
||||
|
||||
# Flat-style implicit scan (only when no nested models: block)
|
||||
if block_base and not has_models_block:
|
||||
_implicit_scan(block_base, flat_model_keys, block_is_default)
|
||||
|
||||
Loading…
Reference in New Issue
Block a user