Harden some nodes against potential issues related to combos. (#15277)
This commit is contained in:
parent
16e3f3034f
commit
9a9fdb10ed
|
|
@ -305,6 +305,12 @@
|
|||
|
||||
- Follow existing node conventions: `INPUT_TYPES`, `RETURN_TYPES`, `FUNCTION`,
|
||||
`CATEGORY`, and registration through the local mapping used by that file.
|
||||
- Treat legacy combo inputs, `io.Combo`, and `io.DynamicCombo` values as
|
||||
untrusted when they affect filesystem access. Any value used as a file or
|
||||
folder name, path component, format, or extension must be validated again at
|
||||
the load/save boundary using an existing `folder_paths` resolver or
|
||||
containment helper, or a fixed allowlist/mapping. Do not rely only on the
|
||||
advertised combo options or prompt validation.
|
||||
- Keep node changes backward compatible by default. Add inputs with sensible
|
||||
defaults and avoid changing output types unless the request requires it.
|
||||
- Model implementations should add the minimal number of ComfyUI nodes required
|
||||
|
|
|
|||
|
|
@ -260,6 +260,7 @@ class ImageSaveHelper:
|
|||
class AudioSaveHelper:
|
||||
"""A helper class with static methods to handle audio saving and metadata."""
|
||||
_OPUS_RATES = [8000, 12000, 16000, 24000, 48000]
|
||||
_FORMATS = {"flac", "mp3", "opus"}
|
||||
|
||||
@staticmethod
|
||||
def save_audio(
|
||||
|
|
@ -270,6 +271,9 @@ class AudioSaveHelper:
|
|||
format: str = "flac",
|
||||
quality: str = "128k",
|
||||
) -> list[SavedResult]:
|
||||
if format not in AudioSaveHelper._FORMATS:
|
||||
raise ValueError(f"Unsupported audio format: {format!r}")
|
||||
|
||||
full_output_folder, filename, counter, subfolder, _ = folder_paths.get_save_image_path(
|
||||
filename_prefix, _get_directory_by_folder_type(folder_type)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -195,7 +195,7 @@ class LoadImageDataSetFromFolderNode(io.ComfyNode):
|
|||
|
||||
@classmethod
|
||||
def execute(cls, folder):
|
||||
sub_input_dir = os.path.join(folder_paths.get_input_directory(), folder)
|
||||
sub_input_dir = secure_subfolder_path(folder_paths.get_input_directory(), folder)
|
||||
valid_extensions = [".png", ".jpg", ".jpeg", ".webp"]
|
||||
image_files = [
|
||||
f
|
||||
|
|
@ -241,7 +241,7 @@ class LoadImageTextDataSetFromFolderNode(io.ComfyNode):
|
|||
def execute(cls, folder):
|
||||
logging.info(f"Loading images from folder: {folder}")
|
||||
|
||||
sub_input_dir = os.path.join(folder_paths.get_input_directory(), folder)
|
||||
sub_input_dir = secure_subfolder_path(folder_paths.get_input_directory(), folder)
|
||||
valid_extensions = [".png", ".jpg", ".jpeg", ".webp"]
|
||||
|
||||
image_files = []
|
||||
|
|
@ -310,7 +310,7 @@ class LoadVideoDataSetFromFolderNode(io.ComfyNode):
|
|||
|
||||
@classmethod
|
||||
def execute(cls, folder):
|
||||
sub_input_dir = os.path.join(folder_paths.get_input_directory(), folder)
|
||||
sub_input_dir = secure_subfolder_path(folder_paths.get_input_directory(), folder)
|
||||
video_files = sorted([
|
||||
f for f in os.listdir(sub_input_dir)
|
||||
if any(f.lower().endswith(ext) for ext in VALID_VIDEO_EXTENSIONS)
|
||||
|
|
@ -357,7 +357,7 @@ class LoadVideoTextDataSetFromFolderNode(io.ComfyNode):
|
|||
|
||||
@classmethod
|
||||
def execute(cls, folder):
|
||||
sub_input_dir = os.path.join(folder_paths.get_input_directory(), folder)
|
||||
sub_input_dir = secure_subfolder_path(folder_paths.get_input_directory(), folder)
|
||||
|
||||
video_files = []
|
||||
for item in sorted(os.listdir(sub_input_dir)):
|
||||
|
|
|
|||
|
|
@ -471,6 +471,12 @@ def _mat_to_quat(m):
|
|||
|
||||
|
||||
class SplatToFile3D(IO.ComfyNode):
|
||||
FORMAT_WRITERS = {
|
||||
"ply": _gaussian_ply_bytes,
|
||||
"ksplat": _gaussian_ksplat_bytes,
|
||||
"spz": _gaussian_spz_bytes,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
|
|
@ -482,7 +488,7 @@ class SplatToFile3D(IO.ComfyNode):
|
|||
"Supports one item per batch only.",
|
||||
inputs=[
|
||||
IO.Splat.Input("splat"),
|
||||
IO.Combo.Input("format", options=["ply", "ksplat", "spz"], # TODO: add "splat" when we have a writer for it
|
||||
IO.Combo.Input("format", options=list(cls.FORMAT_WRITERS), # TODO: add "splat" when we have a writer for it
|
||||
tooltip="ply: standard 3D Gaussian Splat with full spherical harmonics. "
|
||||
"ksplat: mkkellogg SplatBuffer (level 0, uncompressed), base color only "
|
||||
"spz: Niantic gzip-compressed (~10x smaller), base color only "
|
||||
|
|
@ -493,10 +499,13 @@ class SplatToFile3D(IO.ComfyNode):
|
|||
|
||||
@classmethod
|
||||
def execute(cls, splat, format="ply") -> IO.NodeOutput:
|
||||
writer = cls.FORMAT_WRITERS.get(format)
|
||||
if writer is None:
|
||||
raise ValueError(f"Unsupported splat format: {format!r}")
|
||||
|
||||
if splat.positions.shape[0] > 1:
|
||||
logging.warning("SplatToFile3D supports one item per batch only. Got %d; using first.", splat.positions.shape[0])
|
||||
end = _real_len(splat, 0)
|
||||
writer = {"ksplat": _gaussian_ksplat_bytes, "spz": _gaussian_spz_bytes}.get(format, _gaussian_ply_bytes)
|
||||
data = writer(splat.positions[0, :end], splat.scales[0, :end],
|
||||
splat.rotations[0, :end], splat.opacities[0, :end], splat.sh[0, :end])
|
||||
return IO.NodeOutput(Types.File3D(BytesIO(data), file_format=format))
|
||||
|
|
|
|||
19
nodes.py
19
nodes.py
|
|
@ -633,15 +633,18 @@ class DiffusersLoader:
|
|||
SEARCH_ALIASES = ["load diffusers model"]
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
def _model_paths(cls):
|
||||
paths = []
|
||||
for search_path in folder_paths.get_folder_paths("diffusers"):
|
||||
if os.path.exists(search_path):
|
||||
for root, subdir, files in os.walk(search_path, followlinks=True):
|
||||
if "model_index.json" in files:
|
||||
paths.append(os.path.relpath(root, start=search_path))
|
||||
return paths
|
||||
|
||||
return {"required": {"model_path": (paths,), }}
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"model_path": (cls._model_paths(),), }}
|
||||
RETURN_TYPES = ("MODEL", "CLIP", "VAE")
|
||||
FUNCTION = "load_checkpoint"
|
||||
DEPRECATED = True
|
||||
|
|
@ -649,14 +652,20 @@ class DiffusersLoader:
|
|||
CATEGORY = "model/loaders"
|
||||
|
||||
def load_checkpoint(self, model_path, output_vae=True, output_clip=True):
|
||||
if model_path not in self._model_paths():
|
||||
raise ValueError(f"Invalid diffusers model path: {model_path!r}")
|
||||
|
||||
resolved_model_path = None
|
||||
for search_path in folder_paths.get_folder_paths("diffusers"):
|
||||
if os.path.exists(search_path):
|
||||
path = os.path.join(search_path, model_path)
|
||||
if os.path.exists(path):
|
||||
model_path = path
|
||||
if os.path.isfile(os.path.join(path, "model_index.json")):
|
||||
resolved_model_path = path
|
||||
break
|
||||
if resolved_model_path is None:
|
||||
raise FileNotFoundError(f"Diffusers model {model_path!r} not found.")
|
||||
|
||||
return comfy.diffusers_load.load_diffusers(model_path, output_vae=output_vae, output_clip=output_clip, embedding_directory=folder_paths.get_folder_paths("embeddings"))
|
||||
return comfy.diffusers_load.load_diffusers(resolved_model_path, output_vae=output_vae, output_clip=output_clip, embedding_directory=folder_paths.get_folder_paths("embeddings"))
|
||||
|
||||
|
||||
class unCLIPCheckpointLoader:
|
||||
|
|
|
|||
Loading…
Reference in New Issue