Compare commits

...

11 Commits

Author SHA1 Message Date
Nicolò Paternoster
2f54dd88cc
Merge 15d49a61b8 into 6b61918a16 2026-05-19 09:27:05 +00:00
adv0r
15d49a61b8 Address review feedback on /internal/models/download
- Disable aiohttp auto-redirects and re-validate every Location target
  against the same allowlist used for the initial URL, closing an SSRF
  vector where an allowed host could redirect to an arbitrary internal
  endpoint.
- Accept subdomains of allowlisted hosts so Hugging Face's LFS CDN
  (cdn-lfs.huggingface.co et al.) keeps working under the stricter
  redirect handling.
- Pass an explicit ClientTimeout (connect/sock_read) so hung remotes
  surface as errors instead of blocking the request handler forever.
- Log the exception value alongside the traceback on the 500 fallback.
- Add positive coverage for normalize_model_relative_path, Civitai URL
  allowlisting, and the redirect-following / SSRF-rejection branches of
  open_model_download_response.

Co-authored-by: Cursor <cursoragent@cursor.com>
2026-05-19 11:26:53 +02:00
Matt Miller
6b61918a16
docs(openapi): deprecate /api/upload/mask in favor of /api/upload/image (#13968)
Some checks are pending
Python Linting / Run Ruff (push) Waiting to run
Python Linting / Run Pylint (push) Waiting to run
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.10, [self-hosted Linux], stable) (push) Waiting to run
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.11, [self-hosted Linux], stable) (push) Waiting to run
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.12, [self-hosted Linux], stable) (push) Waiting to run
Full Comfy CI Workflow Runs / test-unix-nightly (12.1, , linux, 3.11, [self-hosted Linux], nightly) (push) Waiting to run
Execution Tests / test (macos-latest) (push) Waiting to run
Execution Tests / test (ubuntu-latest) (push) Waiting to run
Execution Tests / test (windows-latest) (push) Waiting to run
Test server launches without errors / test (push) Waiting to run
Unit Tests / test (macos-latest) (push) Waiting to run
Unit Tests / test (ubuntu-latest) (push) Waiting to run
Unit Tests / test (windows-2022) (push) Waiting to run
Mark the uploadMask operation as deprecated and point clients at
/api/upload/image. The mask-compositing behavior the endpoint provides
(alpha-compositing the supplied mask onto an original_ref image) is now
expected to happen client-side, with the composited result uploaded
through the unified /api/upload/image path.

The endpoint continues to function for older clients; no runtime
behavior changes ship with this commit. Only the OpenAPI annotation
and the human-facing description are updated.
2026-05-19 12:19:51 +08:00
comfyanonymous
a4382e056e
Use temporal downscale to make empty audio latent nodes more reusable. (#13975) 2026-05-19 00:14:30 -04:00
Alexis Rolland
d71cc1c8f2
chore: Various QoL updates of nodes display names, descriptions and categories (CORE-190, CORE-191) (#13830)
* Move detection category under image category

* Add missing categories

* Move detection nodes to detection category

* Move save nodes to image root catefory

* Rename postprocessors

* Move mask category under image

* Move guiders category to parent level at root of sampling category

* Move custom_sampling category to parent level at the root of sampling category

* Modify description of LoRA loaders

* Fix node id SolidMask

* Move VOID Quadmask under image/mask

* Group compositing nodes under image/compositing

* Move load image as mask to image category for consistency with other load image nodes

* Align display name with Load Checkpoint

* Move dataset category under training category

* Rename Number Convert to Conver Number (verb first)

* Rename Canny node

* Revert wanBlockSwap + description

* Add description to RemoveBackground node

* Revert category update of dataset
2026-05-19 00:13:48 -04:00
comfyanonymous
990a7ae7f2
Initial work to make downscale_ratio_temporal work. (#13972) 2026-05-18 23:01:43 -04:00
Jedrzej Kosinski
df2454b47e
Reduce min for Batch Image/Mask/Latent nodes from 2 to 1 (#13721) 2026-05-19 09:50:14 +08:00
drozbay
292814c31e
feat: Add optional attention_mask input to LTXVAddGuide (CORE-220) (#13965)
Some checks are pending
Python Linting / Run Ruff (push) Waiting to run
Python Linting / Run Pylint (push) Waiting to run
Execution Tests / test (ubuntu-latest) (push) Waiting to run
Execution Tests / test (windows-latest) (push) Waiting to run
Unit Tests / test (macos-latest) (push) Waiting to run
Unit Tests / test (ubuntu-latest) (push) Waiting to run
Unit Tests / test (windows-2022) (push) Waiting to run
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.10, [self-hosted Linux], stable) (push) Waiting to run
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.11, [self-hosted Linux], stable) (push) Waiting to run
Full Comfy CI Workflow Runs / test-stable (12.1, , linux, 3.12, [self-hosted Linux], stable) (push) Waiting to run
Full Comfy CI Workflow Runs / test-unix-nightly (12.1, , linux, 3.11, [self-hosted Linux], nightly) (push) Waiting to run
Execution Tests / test (macos-latest) (push) Waiting to run
Test server launches without errors / test (push) Waiting to run
2026-05-19 05:07:04 +08:00
Yousef R. Gamaleldin
187e5237e1
Fix BiRefNet issue (#13966) 2026-05-19 05:03:22 +08:00
Alexander Piskun
164a9d4bbb
[Partner Nodes] add ByteDance Seed LLM node (#13919)
Signed-off-by: bigcat88 <bigcat88@icloud.com>
2026-05-18 13:06:13 -07:00
rattus
16f862f02a
implement dynamic clip saving (#13959)
Fix clip saving by doing the same patching process and diffusion
models.
2026-05-18 11:46:40 -07:00
37 changed files with 753 additions and 129 deletions

View File

@ -72,8 +72,8 @@ class InternalRoutes:
)
except ModelDownloadError as err:
return web.json_response({"error": str(err)}, status=err.status)
except Exception:
logging.exception("Failed to download model")
except Exception as err:
logging.exception("Failed to download model: %s", err)
return web.json_response({"error": "Failed to download model."}, status=500)
response_status = 200 if result["status"] == "already_exists" else 201

View File

@ -4,9 +4,9 @@ import os
import posixpath
from dataclasses import dataclass
from pathlib import PurePosixPath
from urllib.parse import urlparse
from urllib.parse import urljoin, urlparse
from aiohttp import ClientSession
from aiohttp import ClientSession, ClientTimeout
import folder_paths
@ -16,6 +16,18 @@ ALLOWED_DOWNLOAD_SUFFIXES = (".safetensors", ".sft", ".ckpt", ".pth", ".pt")
BLOCKED_MODEL_FOLDERS = {"configs", "custom_nodes"}
CHUNK_SIZE = 1024 * 1024
# Bound the network call so a hung remote eventually surfaces an error
# instead of blocking the request handler forever. ``sock_read`` is the
# inter-chunk read timeout, which is the right knob for long downloads:
# a slow-but-progressing transfer keeps making progress, while a stalled
# socket fails predictably.
DOWNLOAD_TIMEOUT = ClientTimeout(total=None, connect=30, sock_connect=30, sock_read=300)
# Maximum number of redirects we follow manually. Hugging Face typically
# redirects ``/resolve/main/...`` to a single CDN URL, so a small budget
# is enough while still preventing redirect loops.
MAX_DOWNLOAD_REDIRECTS = 5
WHITE_LISTED_DOWNLOAD_URLS = {
"https://huggingface.co/stabilityai/stable-zero123/resolve/main/stable_zero123.ckpt",
"https://huggingface.co/TencentARC/T2I-Adapter/resolve/main/models/t2iadapter_depth_sd14v1.pth?download=true",
@ -71,6 +83,14 @@ def parse_model_download_request(data) -> ModelDownloadRequest:
def is_allowed_model_download_url(url: str) -> bool:
"""Return True for URLs we are willing to fetch on behalf of the user.
The same predicate is applied to the user-supplied URL and to every
redirect target, so SSRF via redirects on an allowed host is contained
to the same allowlist. Subdomains of allowlisted hosts are accepted
because Hugging Face and Civitai both serve actual file payloads from
CDN subdomains (e.g. ``cdn-lfs.huggingface.co``).
"""
if url in WHITE_LISTED_DOWNLOAD_URLS:
return True
@ -82,7 +102,14 @@ def is_allowed_model_download_url(url: str) -> bool:
if parsed.scheme != "https":
return False
return (parsed.hostname or "").lower() in ALLOWED_DOWNLOAD_HOSTS
host = (parsed.hostname or "").lower()
if not host:
return False
for allowed in ALLOWED_DOWNLOAD_HOSTS:
if host == allowed or host.endswith("." + allowed):
return True
return False
def normalize_model_relative_path(name: str) -> str:
@ -164,6 +191,36 @@ def safe_join(root: str, relative_path: str) -> str:
return full_path
async def open_model_download_response(session: ClientSession, url: str):
"""GET ``url`` with explicit timeout and an allowlist-checked redirect chain.
aiohttp follows redirects by default, which would let an allowed host
redirect to an arbitrary internal target (SSRF). We disable automatic
following and validate every ``Location`` against the same allowlist
used for the initial URL.
"""
current_url = url
for _ in range(MAX_DOWNLOAD_REDIRECTS + 1):
response = await session.get(
current_url,
allow_redirects=False,
timeout=DOWNLOAD_TIMEOUT,
)
if response.status not in (301, 302, 303, 307, 308):
return response
location = response.headers.get("Location", "").strip()
response.release()
if not location:
raise ModelDownloadError("Redirect response missing Location header.", status=502)
next_url = urljoin(current_url, location)
if not is_allowed_model_download_url(next_url):
raise ModelDownloadError("Model download redirect target is not allowed.", status=403)
current_url = next_url
raise ModelDownloadError("Too many redirects while downloading model.", status=502)
async def download_model_to_destination(
session: ClientSession,
request: ModelDownloadRequest,
@ -183,7 +240,7 @@ async def download_model_to_destination(
bytes_written = 0
try:
with os.fdopen(fd, "wb") as output:
async with session.get(request.url) as response:
async with await open_model_download_response(session, request.url) as response:
if response.status >= 400:
raise ModelDownloadError(f"Model download failed with HTTP {response.status}.", status=502)

View File

@ -44,7 +44,14 @@ class BackgroundRemovalModel():
comfy.model_management.load_model_gpu(self.patcher)
H, W = image.shape[1], image.shape[2]
pixel_values = comfy.clip_model.clip_preprocess(image.to(self.load_device), size=self.image_size, mean=self.image_mean, std=self.image_std, crop=False)
out = self.model(pixel_values=pixel_values)
if pixel_values.shape[0] > 1:
out = torch.cat([
self.model(pixel_values=pixel_values[i:i+1])
for i in range(pixel_values.shape[0])
], dim=0)
else:
out = self.model(pixel_values=pixel_values)
out = torch.nn.functional.interpolate(out, size=(H, W), mode="bicubic", antialias=False)
mask = out.sigmoid().to(device=comfy.model_management.intermediate_device(), dtype=comfy.model_management.intermediate_dtype())

View File

@ -150,6 +150,7 @@ class SD3(LatentFormat):
class StableAudio1(LatentFormat):
latent_channels = 64
latent_dimensions = 1
temporal_downscale_ratio = 2048
class Flux(SD3):
latent_channels = 16
@ -766,6 +767,7 @@ class ACEAudio(LatentFormat):
class ACEAudio15(LatentFormat):
latent_channels = 64
latent_dimensions = 1
temporal_downscale_ratio = 1764
class ChromaRadiance(LatentFormat):
latent_channels = 3

View File

@ -1493,27 +1493,30 @@ class ModelPatcher:
self.unpatch_hooks()
self.clear_cached_hook_weights()
def state_dict_for_saving(self, clip_state_dict=None, vae_state_dict=None, clip_vision_state_dict=None):
original_state_dict = self.model.diffusion_model.state_dict()
unet_state_dict = {}
def model_state_dict_for_saving(self, model=None, prefix=""):
if model is None:
model = self.model
original_state_dict = model.state_dict()
output_state_dict = {}
keys = list(original_state_dict)
while len(keys) > 0:
k = keys.pop(0)
v = original_state_dict[k]
op_keys = k.rsplit('.', 1)
if (len(op_keys) < 2) or op_keys[1] not in ["weight", "bias"]:
unet_state_dict[k] = v
output_state_dict[k] = v
continue
try:
op = comfy.utils.get_attr(self.model.diffusion_model, op_keys[0])
op = comfy.utils.get_attr(model, op_keys[0])
except:
unet_state_dict[k] = v
output_state_dict[k] = v
continue
if not op or not hasattr(op, "comfy_cast_weights") or \
(hasattr(op, "comfy_patched_weights") and op.comfy_patched_weights == True):
unet_state_dict[k] = v
output_state_dict[k] = v
continue
key = "diffusion_model." + k
key = prefix + k
weight = comfy.utils.get_attr(self.model, key)
if isinstance(weight, QuantizedTensor) and k in original_state_dict:
qt_state_dict = weight.state_dict(k)
@ -1521,10 +1524,14 @@ class ModelPatcher:
for group_key in (x for x in qt_state_dict if x in original_state_dict):
if group_key in keys:
keys.remove(group_key)
unet_state_dict.pop(group_key, "")
unet_state_dict[group_key] = LazyCastingParamPiece(caster, "diffusion_model." + group_key, original_state_dict[group_key])
output_state_dict.pop(group_key, "")
output_state_dict[group_key] = LazyCastingParamPiece(caster, prefix + group_key, original_state_dict[group_key])
continue
unet_state_dict[k] = LazyCastingParam(self, key, weight)
output_state_dict[k] = LazyCastingParam(self, key, weight)
return output_state_dict
def state_dict_for_saving(self, clip_state_dict=None, vae_state_dict=None, clip_vision_state_dict=None):
unet_state_dict = self.model_state_dict_for_saving(self.model.diffusion_model, "diffusion_model.")
return self.model.state_dict_for_saving(unet_state_dict, clip_state_dict=clip_state_dict, vae_state_dict=vae_state_dict, clip_vision_state_dict=clip_vision_state_dict)
def __del__(self):

View File

@ -37,11 +37,12 @@ def prepare_noise(latent_image, seed, noise_inds=None):
return noises
def fix_empty_latent_channels(model, latent_image, downscale_ratio_spacial=None):
def fix_empty_latent_channels(model, latent_image, downscale_ratio_spacial=None, downscale_ratio_temporal=None):
if latent_image.is_nested:
return latent_image
latent_format = model.get_model_object("latent_format") #Resize the empty latent image so it has the right number of channels
if torch.count_nonzero(latent_image) == 0:
is_empty = torch.count_nonzero(latent_image) == 0
if is_empty:
if latent_format.latent_channels != latent_image.shape[1]:
latent_image = comfy.utils.repeat_to_batch_size(latent_image, latent_format.latent_channels, dim=1)
if downscale_ratio_spacial is not None:
@ -51,6 +52,13 @@ def fix_empty_latent_channels(model, latent_image, downscale_ratio_spacial=None)
if latent_format.latent_dimensions == 3 and latent_image.ndim == 4:
latent_image = latent_image.unsqueeze(2)
if is_empty and downscale_ratio_temporal is not None:
if downscale_ratio_temporal != latent_format.temporal_downscale_ratio:
ratio = downscale_ratio_temporal / latent_format.temporal_downscale_ratio
new_t = max(1, round(latent_image.shape[2] * ratio))
latent_image = comfy.utils.repeat_to_batch_size(latent_image, new_t, dim=2)
return latent_image
def prepare_sampling(model, noise_shape, positive, negative, noise_mask):

View File

@ -423,6 +423,13 @@ class CLIP:
sd_clip[k] = sd_tokenizer[k]
return sd_clip
def state_dict_for_saving(self):
sd_clip = self.patcher.model_state_dict_for_saving()
sd_tokenizer = self.tokenizer.state_dict()
for k in sd_tokenizer:
sd_clip[k] = sd_tokenizer[k]
return sd_clip
def load_model(self, tokens={}):
memory_used = 0
if hasattr(self.cond_stage_model, "memory_estimation_function"):
@ -1908,7 +1915,7 @@ def save_checkpoint(output_path, model, clip=None, vae=None, clip_vision=None, m
load_models = [model]
if clip is not None:
load_models.append(clip.load_model())
clip_sd = clip.get_sd()
clip_sd = clip.state_dict_for_saving()
vae_sd = None
if vae is not None:
vae_sd = vae.get_sd()

View File

@ -0,0 +1,101 @@
"""Pydantic models for BytePlus ModelArk Responses API.
See: https://docs.byteplus.com/en/docs/ModelArk/1585128 (request)
https://docs.byteplus.com/en/docs/ModelArk/1783703 (response)
"""
from typing import Literal
from pydantic import BaseModel, Field
class BytePlusInputText(BaseModel):
type: Literal["input_text"] = "input_text"
text: str = Field(...)
class BytePlusInputImage(BaseModel):
type: Literal["input_image"] = "input_image"
image_url: str = Field(..., description="Image URL or `data:image/...;base64,...` payload")
detail: str = Field("auto", description="One of high, low, auto")
class BytePlusInputVideo(BaseModel):
type: Literal["input_video"] = "input_video"
video_url: str = Field(..., description="Video URL or `data:video/...;base64,...` payload")
fps: float | None = Field(None, ge=0.2, le=5.0)
BytePlusMessageContent = BytePlusInputText | BytePlusInputImage | BytePlusInputVideo
class BytePlusInputMessage(BaseModel):
type: Literal["message"] = "message"
role: str = Field(..., description="One of user, system, assistant, developer")
content: list[BytePlusMessageContent] = Field(...)
class BytePlusResponseCreateRequest(BaseModel):
model: str = Field(...)
input: list[BytePlusInputMessage] = Field(...)
instructions: str | None = Field(None)
max_output_tokens: int | None = Field(None, ge=1)
temperature: float | None = Field(None, ge=0.0, le=2.0)
store: bool | None = Field(False)
stream: bool | None = Field(False)
class BytePlusOutputText(BaseModel):
type: Literal["output_text"] = "output_text"
text: str = Field(...)
class BytePlusOutputRefusal(BaseModel):
type: Literal["refusal"] = "refusal"
refusal: str = Field(...)
class BytePlusOutputContent(BaseModel):
type: str = Field(...)
text: str | None = Field(None)
refusal: str | None = Field(None)
class BytePlusOutputMessage(BaseModel):
type: str = Field(...)
id: str | None = Field(None)
role: str | None = Field(None)
status: str | None = Field(None)
content: list[BytePlusOutputContent] | None = Field(None)
class BytePlusInputTokensDetails(BaseModel):
cached_tokens: int | None = Field(None)
class BytePlusOutputTokensDetails(BaseModel):
reasoning_tokens: int | None = Field(None)
class BytePlusResponseUsage(BaseModel):
input_tokens: int | None = Field(None)
output_tokens: int | None = Field(None)
total_tokens: int | None = Field(None)
input_tokens_details: BytePlusInputTokensDetails | None = Field(None)
output_tokens_details: BytePlusOutputTokensDetails | None = Field(None)
class BytePlusResponseError(BaseModel):
code: str = Field(...)
message: str = Field(...)
class BytePlusResponseObject(BaseModel):
id: str | None = Field(None)
object: str | None = Field(None)
created_at: int | None = Field(None)
model: str | None = Field(None)
status: str | None = Field(None)
error: BytePlusResponseError | None = Field(None)
output: list[BytePlusOutputMessage] | None = Field(None)
usage: BytePlusResponseUsage | None = Field(None)

View File

@ -0,0 +1,271 @@
"""API Nodes for ByteDance Seed LLM via the BytePlus ModelArk Responses API.
See: https://docs.byteplus.com/en/docs/ModelArk/1585128
"""
from typing_extensions import override
from comfy_api.latest import IO, ComfyExtension, Input
from comfy_api_nodes.apis.bytedance_llm import (
BytePlusInputImage,
BytePlusInputMessage,
BytePlusInputText,
BytePlusInputVideo,
BytePlusMessageContent,
BytePlusResponseCreateRequest,
BytePlusResponseObject,
)
from comfy_api_nodes.util import (
ApiEndpoint,
get_number_of_images,
sync_op,
upload_images_to_comfyapi,
upload_video_to_comfyapi,
validate_string,
)
BYTEPLUS_RESPONSES_ENDPOINT = "/proxy/byteplus/api/v3/responses"
SEED_MAX_IMAGES = 20
SEED_MAX_VIDEOS = 4
SEED_MODELS: dict[str, str] = {
"Seed 2.0 Pro": "seed-2-0-pro-260328",
"Seed 2.0 Lite": "seed-2-0-lite-260228",
"Seed 2.0 Mini": "seed-2-0-mini-260215",
}
# USD per 1M tokens: (input, cache_hit_input, output)
_SEED_PRICES_PER_MILLION: dict[str, tuple[float, float, float]] = {
"seed-2-0-pro-260328": (0.50, 0.10, 3.00),
"seed-2-0-lite-260228": (0.25, 0.05, 2.00),
"seed-2-0-mini-260215": (0.10, 0.02, 0.40),
}
def _seed_model_inputs(max_images: int = SEED_MAX_IMAGES, max_videos: int = SEED_MAX_VIDEOS):
return [
IO.Autogrow.Input(
"images",
template=IO.Autogrow.TemplateNames(
IO.Image.Input("image"),
names=[f"image_{i}" for i in range(1, max_images + 1)],
min=0,
),
tooltip=f"Optional image(s) to use as context for the model. Up to {max_images} images.",
),
IO.Autogrow.Input(
"videos",
template=IO.Autogrow.TemplateNames(
IO.Video.Input("video"),
names=[f"video_{i}" for i in range(1, max_videos + 1)],
min=0,
),
tooltip=f"Optional video(s) to use as context for the model. Up to {max_videos} videos.",
),
IO.Float.Input(
"temperature",
default=1.0,
min=0.0,
max=2.0,
step=0.01,
tooltip="Controls randomness. 0.0 is deterministic, higher values are more random.",
advanced=True,
),
]
def _calculate_price(model_id: str, response: BytePlusResponseObject) -> float | None:
"""Compute approximate USD price from response usage."""
if not response.usage:
return None
rates = _SEED_PRICES_PER_MILLION.get(model_id)
if rates is None:
return None
input_rate, cache_hit_rate, output_rate = rates
input_tokens = response.usage.input_tokens or 0
output_tokens = response.usage.output_tokens or 0
cached = 0
if response.usage.input_tokens_details:
cached = response.usage.input_tokens_details.cached_tokens or 0
fresh_input = max(0, input_tokens - cached)
total = fresh_input * input_rate + cached * cache_hit_rate + output_tokens * output_rate
return total / 1_000_000.0
def _get_text_from_response(response: BytePlusResponseObject) -> str:
"""Extract concatenated text from all assistant message output_text blocks."""
if not response.output:
return ""
chunks: list[str] = []
for item in response.output:
if item.type != "message" or not item.content:
continue
for block in item.content:
if block.type == "output_text" and block.text:
chunks.append(block.text)
elif block.type == "refusal" and block.refusal:
raise ValueError(f"Model refused to respond: {block.refusal}")
return "\n".join(chunks)
async def _build_image_content_blocks(
cls: type[IO.ComfyNode],
image_tensors: list[Input.Image],
) -> list[BytePlusInputImage]:
urls = await upload_images_to_comfyapi(
cls,
image_tensors,
max_images=SEED_MAX_IMAGES,
wait_label="Uploading reference images",
)
return [BytePlusInputImage(image_url=url) for url in urls]
async def _build_video_content_blocks(
cls: type[IO.ComfyNode],
videos: list[Input.Video],
) -> list[BytePlusInputVideo]:
blocks: list[BytePlusInputVideo] = []
total = len(videos)
for idx, video in enumerate(videos):
label = "Uploading reference video"
if total > 1:
label = f"{label} ({idx + 1}/{total})"
url = await upload_video_to_comfyapi(cls, video, wait_label=label)
blocks.append(BytePlusInputVideo(video_url=url))
return blocks
class ByteDanceSeedNode(IO.ComfyNode):
"""Generate text responses from a ByteDance Seed 2.0 model."""
@classmethod
def define_schema(cls):
return IO.Schema(
node_id="ByteDanceSeedNode",
display_name="ByteDance Seed",
category="api node/text/ByteDance",
essentials_category="Text Generation",
description="Generate text responses with ByteDance's Seed 2.0 models. "
"Provide a text prompt and optionally one or more images or videos for multimodal context.",
inputs=[
IO.String.Input(
"prompt",
multiline=True,
default="",
tooltip="Text input to the model.",
),
IO.DynamicCombo.Input(
"model",
options=[IO.DynamicCombo.Option(label, _seed_model_inputs()) for label in SEED_MODELS],
tooltip="The Seed model used to generate the response.",
),
IO.Int.Input(
"seed",
default=0,
min=0,
max=2147483647,
control_after_generate=True,
tooltip="Seed controls whether the node should re-run; "
"results are non-deterministic regardless of seed.",
),
IO.String.Input(
"system_prompt",
multiline=True,
default="",
optional=True,
advanced=True,
tooltip="Foundational instructions that dictate the model's behavior.",
),
],
outputs=[IO.String.Output()],
hidden=[
IO.Hidden.auth_token_comfy_org,
IO.Hidden.api_key_comfy_org,
IO.Hidden.unique_id,
],
is_api_node=True,
price_badge=IO.PriceBadge(
depends_on=IO.PriceBadgeDepends(widgets=["model"]),
expr="""
(
$m := widgets.model;
$contains($m, "mini") ? {
"type": "list_usd",
"usd": [0.00025, 0.0009],
"format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
}
: $contains($m, "lite") ? {
"type": "list_usd",
"usd": [0.0003, 0.002],
"format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
}
: $contains($m, "pro") ? {
"type": "list_usd",
"usd": [0.0005, 0.003],
"format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
}
: {"type":"text", "text":"Token-based"}
)
""",
),
)
@classmethod
async def execute(
cls,
prompt: str,
model: dict,
seed: int,
system_prompt: str = "",
) -> IO.NodeOutput:
validate_string(prompt, strip_whitespace=True, min_length=1)
model_label = model["model"]
temperature = model["temperature"]
model_id = SEED_MODELS[model_label]
image_tensors: list[Input.Image] = [t for t in (model.get("images") or {}).values() if t is not None]
if sum(get_number_of_images(t) for t in image_tensors) > SEED_MAX_IMAGES:
raise ValueError(f"Up to {SEED_MAX_IMAGES} images are supported per request.")
video_inputs: list[Input.Video] = [v for v in (model.get("videos") or {}).values() if v is not None]
if len(video_inputs) > SEED_MAX_VIDEOS:
raise ValueError(f"Up to {SEED_MAX_VIDEOS} videos are supported per request.")
content: list[BytePlusMessageContent] = []
if image_tensors:
content.extend(await _build_image_content_blocks(cls, image_tensors))
if video_inputs:
content.extend(await _build_video_content_blocks(cls, video_inputs))
content.append(BytePlusInputText(text=prompt))
response = await sync_op(
cls,
ApiEndpoint(path=BYTEPLUS_RESPONSES_ENDPOINT, method="POST"),
response_model=BytePlusResponseObject,
data=BytePlusResponseCreateRequest(
model=model_id,
input=[BytePlusInputMessage(role="user", content=content)],
instructions=system_prompt or None,
temperature=temperature,
store=False,
stream=False,
),
price_extractor=lambda r: _calculate_price(model_id, r),
)
if response.error:
raise ValueError(f"Seed API error ({response.error.code}): {response.error.message}")
result = _get_text_from_response(response)
if not result:
raise ValueError("Empty response from Seed model.")
return IO.NodeOutput(result)
class ByteDanceLLMExtension(ComfyExtension):
@override
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
return [ByteDanceSeedNode]
async def comfy_entrypoint() -> ByteDanceLLMExtension:
return ByteDanceLLMExtension()

View File

@ -104,7 +104,7 @@ class EmptyAceStep15LatentAudio(IO.ComfyNode):
def execute(cls, seconds, batch_size) -> IO.NodeOutput:
length = round((seconds * 48000 / 1920))
latent = torch.zeros([batch_size, 64, length], device=comfy.model_management.intermediate_device(), dtype=comfy.model_management.intermediate_dtype())
return IO.NodeOutput({"samples": latent, "type": "audio"})
return IO.NodeOutput({"samples": latent, "type": "audio", "downscale_ratio_temporal": 1764})
class ReferenceAudio(IO.ComfyNode):
@classmethod

View File

@ -45,7 +45,7 @@ class SamplerLCMUpscale(io.ComfyNode):
def define_schema(cls) -> io.Schema:
return io.Schema(
node_id="SamplerLCMUpscale",
category="sampling/custom_sampling/samplers",
category="sampling/samplers",
inputs=[
io.Float.Input("scale_ratio", default=1.0, min=0.1, max=20.0, step=0.01, advanced=True),
io.Int.Input("scale_steps", default=-1, min=-1, max=1000, step=1, advanced=True),
@ -123,7 +123,7 @@ class SamplerEulerCFGpp(io.ComfyNode):
return io.Schema(
node_id="SamplerEulerCFGpp",
display_name="SamplerEulerCFG++",
category="experimental", # "sampling/custom_sampling/samplers"
category="experimental", # "sampling/samplers"
inputs=[
io.Combo.Input("version", options=["regular", "alternative"], advanced=True),
],

View File

@ -29,7 +29,7 @@ class AlignYourStepsScheduler(io.ComfyNode):
return io.Schema(
node_id="AlignYourStepsScheduler",
search_aliases=["AYS scheduler"],
category="sampling/custom_sampling/schedulers",
category="sampling/schedulers",
inputs=[
io.Combo.Input("model_type", options=["SD1", "SDXL", "SVD"]),
io.Int.Input("steps", default=10, min=1, max=10000),

View File

@ -53,7 +53,7 @@ class SamplerARVideo(io.ComfyNode):
return io.Schema(
node_id="SamplerARVideo",
display_name="Sampler AR Video",
category="sampling/custom_sampling/samplers",
category="sampling/samplers",
inputs=[
io.Int.Input(
"num_frame_per_block",

View File

@ -33,7 +33,7 @@ class EmptyLatentAudio(IO.ComfyNode):
def execute(cls, seconds, batch_size) -> IO.NodeOutput:
length = round((seconds * 44100 / 2048) / 2) * 2
latent = torch.zeros([batch_size, 64, length], device=comfy.model_management.intermediate_device())
return IO.NodeOutput({"samples":latent, "type": "audio"})
return IO.NodeOutput({"samples": latent, "type": "audio", "downscale_ratio_temporal": 2048})
generate = execute # TODO: remove

View File

@ -34,6 +34,7 @@ class RemoveBackground(IO.ComfyNode):
node_id="RemoveBackground",
display_name="Remove Background",
category="image/background removal",
description="Generates a foreground mask to remove the background from an image using a background removal model.",
inputs=[
IO.Image.Input("image", tooltip="Input image to remove the background from"),
IO.BackgroundRemoval.Input("bg_removal_model", tooltip="Background removal model used to generate the mask")

View File

@ -11,9 +11,9 @@ class Canny(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="Canny",
display_name="Canny",
display_name="Detect Edges (Canny)",
search_aliases=["edge detection", "outline", "contour detection", "line art"],
category="image/preprocessors",
category="image/filters",
essentials_category="Image Tools",
inputs=[
io.Image.Input("image"),

View File

@ -111,7 +111,7 @@ class PorterDuffImageComposite(io.ComfyNode):
node_id="PorterDuffImageComposite",
search_aliases=["alpha composite", "blend modes", "layer blend", "transparency blend"],
display_name="Porter-Duff Image Composite",
category="mask/compositing",
category="image/compositing",
inputs=[
io.Image.Input("source"),
io.Mask.Input("source_alpha"),
@ -168,7 +168,7 @@ class SplitImageWithAlpha(io.ComfyNode):
node_id="SplitImageWithAlpha",
search_aliases=["extract alpha", "separate transparency", "remove alpha"],
display_name="Split Image with Alpha",
category="mask/compositing",
category="image/compositing",
inputs=[
io.Image.Input("image"),
],
@ -192,7 +192,7 @@ class JoinImageWithAlpha(io.ComfyNode):
node_id="JoinImageWithAlpha",
search_aliases=["add transparency", "apply alpha", "composite alpha", "RGBA"],
display_name="Join Image with Alpha",
category="mask/compositing",
category="image/compositing",
inputs=[
io.Image.Input("image"),
io.Mask.Input("alpha"),

View File

@ -17,7 +17,7 @@ class BasicScheduler(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="BasicScheduler",
category="sampling/custom_sampling/schedulers",
category="sampling/schedulers",
inputs=[
io.Model.Input("model"),
io.Combo.Input("scheduler", options=comfy.samplers.SCHEDULER_NAMES),
@ -47,7 +47,7 @@ class KarrasScheduler(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="KarrasScheduler",
category="sampling/custom_sampling/schedulers",
category="sampling/schedulers",
inputs=[
io.Int.Input("steps", default=20, min=1, max=10000),
io.Float.Input("sigma_max", default=14.614642, min=0.0, max=5000.0, step=0.01, round=False, advanced=True),
@ -69,7 +69,7 @@ class ExponentialScheduler(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="ExponentialScheduler",
category="sampling/custom_sampling/schedulers",
category="sampling/schedulers",
inputs=[
io.Int.Input("steps", default=20, min=1, max=10000),
io.Float.Input("sigma_max", default=14.614642, min=0.0, max=5000.0, step=0.01, round=False, advanced=True),
@ -90,7 +90,7 @@ class PolyexponentialScheduler(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="PolyexponentialScheduler",
category="sampling/custom_sampling/schedulers",
category="sampling/schedulers",
inputs=[
io.Int.Input("steps", default=20, min=1, max=10000),
io.Float.Input("sigma_max", default=14.614642, min=0.0, max=5000.0, step=0.01, round=False, advanced=True),
@ -112,7 +112,7 @@ class LaplaceScheduler(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="LaplaceScheduler",
category="sampling/custom_sampling/schedulers",
category="sampling/schedulers",
inputs=[
io.Int.Input("steps", default=20, min=1, max=10000),
io.Float.Input("sigma_max", default=14.614642, min=0.0, max=5000.0, step=0.01, round=False, advanced=True),
@ -136,7 +136,7 @@ class SDTurboScheduler(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="SDTurboScheduler",
category="sampling/custom_sampling/schedulers",
category="sampling/schedulers",
inputs=[
io.Model.Input("model"),
io.Int.Input("steps", default=1, min=1, max=10),
@ -160,7 +160,7 @@ class BetaSamplingScheduler(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="BetaSamplingScheduler",
category="sampling/custom_sampling/schedulers",
category="sampling/schedulers",
inputs=[
io.Model.Input("model"),
io.Int.Input("steps", default=20, min=1, max=10000),
@ -182,7 +182,7 @@ class VPScheduler(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="VPScheduler",
category="sampling/custom_sampling/schedulers",
category="sampling/schedulers",
inputs=[
io.Int.Input("steps", default=20, min=1, max=10000),
io.Float.Input("beta_d", default=19.9, min=0.0, max=5000.0, step=0.01, round=False, advanced=True), #TODO: fix default values
@ -204,7 +204,7 @@ class SplitSigmas(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="SplitSigmas",
category="sampling/custom_sampling/sigmas",
category="sampling/sigmas",
inputs=[
io.Sigmas.Input("sigmas"),
io.Int.Input("step", default=0, min=0, max=10000),
@ -228,7 +228,7 @@ class SplitSigmasDenoise(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="SplitSigmasDenoise",
category="sampling/custom_sampling/sigmas",
category="sampling/sigmas",
inputs=[
io.Sigmas.Input("sigmas"),
io.Float.Input("denoise", default=1.0, min=0.0, max=1.0, step=0.01),
@ -254,7 +254,7 @@ class FlipSigmas(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="FlipSigmas",
category="sampling/custom_sampling/sigmas",
category="sampling/sigmas",
inputs=[io.Sigmas.Input("sigmas")],
outputs=[io.Sigmas.Output()]
)
@ -276,7 +276,7 @@ class SetFirstSigma(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="SetFirstSigma",
category="sampling/custom_sampling/sigmas",
category="sampling/sigmas",
inputs=[
io.Sigmas.Input("sigmas"),
io.Float.Input("sigma", default=136.0, min=0.0, max=20000.0, step=0.001, round=False),
@ -298,7 +298,7 @@ class ExtendIntermediateSigmas(io.ComfyNode):
return io.Schema(
node_id="ExtendIntermediateSigmas",
search_aliases=["interpolate sigmas"],
category="sampling/custom_sampling/sigmas",
category="sampling/sigmas",
inputs=[
io.Sigmas.Input("sigmas"),
io.Int.Input("steps", default=2, min=1, max=100),
@ -351,7 +351,7 @@ class SamplingPercentToSigma(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="SamplingPercentToSigma",
category="sampling/custom_sampling/sigmas",
category="sampling/sigmas",
inputs=[
io.Model.Input("model"),
io.Float.Input("sampling_percent", default=0.0, min=0.0, max=1.0, step=0.0001),
@ -379,7 +379,7 @@ class KSamplerSelect(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="KSamplerSelect",
category="sampling/custom_sampling/samplers",
category="sampling/samplers",
inputs=[io.Combo.Input("sampler_name", options=comfy.samplers.SAMPLER_NAMES)],
outputs=[io.Sampler.Output()]
)
@ -396,7 +396,7 @@ class SamplerDPMPP_3M_SDE(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="SamplerDPMPP_3M_SDE",
category="sampling/custom_sampling/samplers",
category="sampling/samplers",
inputs=[
io.Float.Input("eta", default=1.0, min=0.0, max=100.0, step=0.01, round=False, advanced=True),
io.Float.Input("s_noise", default=1.0, min=0.0, max=100.0, step=0.01, round=False, advanced=True),
@ -421,7 +421,7 @@ class SamplerDPMPP_2M_SDE(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="SamplerDPMPP_2M_SDE",
category="sampling/custom_sampling/samplers",
category="sampling/samplers",
inputs=[
io.Combo.Input("solver_type", options=['midpoint', 'heun']),
io.Float.Input("eta", default=1.0, min=0.0, max=100.0, step=0.01, round=False, advanced=True),
@ -448,7 +448,7 @@ class SamplerDPMPP_SDE(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="SamplerDPMPP_SDE",
category="sampling/custom_sampling/samplers",
category="sampling/samplers",
inputs=[
io.Float.Input("eta", default=1.0, min=0.0, max=100.0, step=0.01, round=False, advanced=True),
io.Float.Input("s_noise", default=1.0, min=0.0, max=100.0, step=0.01, round=False, advanced=True),
@ -474,7 +474,7 @@ class SamplerDPMPP_2S_Ancestral(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="SamplerDPMPP_2S_Ancestral",
category="sampling/custom_sampling/samplers",
category="sampling/samplers",
inputs=[
io.Float.Input("eta", default=1.0, min=0.0, max=100.0, step=0.01, round=False),
io.Float.Input("s_noise", default=1.0, min=0.0, max=100.0, step=0.01, round=False),
@ -494,7 +494,7 @@ class SamplerEulerAncestral(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="SamplerEulerAncestral",
category="sampling/custom_sampling/samplers",
category="sampling/samplers",
inputs=[
io.Float.Input("eta", default=1.0, min=0.0, max=100.0, step=0.01, round=False, advanced=True),
io.Float.Input("s_noise", default=1.0, min=0.0, max=100.0, step=0.01, round=False, advanced=True),
@ -515,7 +515,7 @@ class SamplerEulerAncestralCFGPP(io.ComfyNode):
return io.Schema(
node_id="SamplerEulerAncestralCFGPP",
display_name="SamplerEulerAncestralCFG++",
category="sampling/custom_sampling/samplers",
category="sampling/samplers",
inputs=[
io.Float.Input("eta", default=1.0, min=0.0, max=1.0, step=0.01, round=False),
io.Float.Input("s_noise", default=1.0, min=0.0, max=10.0, step=0.01, round=False),
@ -537,7 +537,7 @@ class SamplerLMS(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="SamplerLMS",
category="sampling/custom_sampling/samplers",
category="sampling/samplers",
inputs=[io.Int.Input("order", default=4, min=1, max=100, advanced=True)],
outputs=[io.Sampler.Output()]
)
@ -554,7 +554,7 @@ class SamplerDPMAdaptative(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="SamplerDPMAdaptative",
category="sampling/custom_sampling/samplers",
category="sampling/samplers",
inputs=[
io.Int.Input("order", default=3, min=2, max=3, advanced=True),
io.Float.Input("rtol", default=0.05, min=0.0, max=100.0, step=0.01, round=False, advanced=True),
@ -585,7 +585,7 @@ class SamplerER_SDE(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="SamplerER_SDE",
category="sampling/custom_sampling/samplers",
category="sampling/samplers",
inputs=[
io.Combo.Input("solver_type", options=["ER-SDE", "Reverse-time SDE", "ODE"]),
io.Int.Input("max_stage", default=3, min=1, max=3, advanced=True),
@ -623,7 +623,7 @@ class SamplerSASolver(io.ComfyNode):
return io.Schema(
node_id="SamplerSASolver",
search_aliases=["sde"],
category="sampling/custom_sampling/samplers",
category="sampling/samplers",
inputs=[
io.Model.Input("model"),
io.Float.Input("eta", default=1.0, min=0.0, max=10.0, step=0.01, round=False, advanced=True),
@ -668,7 +668,7 @@ class SamplerSEEDS2(io.ComfyNode):
return io.Schema(
node_id="SamplerSEEDS2",
search_aliases=["sde", "exp heun"],
category="sampling/custom_sampling/samplers",
category="sampling/samplers",
inputs=[
io.Combo.Input("solver_type", options=["phi_1", "phi_2"]),
io.Float.Input("eta", default=1.0, min=0.0, max=100.0, step=0.01, round=False, tooltip="Stochastic strength", advanced=True),
@ -750,7 +750,7 @@ class SamplerCustom(io.ComfyNode):
latent = latent_image
latent_image = latent["samples"]
latent = latent.copy()
latent_image = comfy.sample.fix_empty_latent_channels(model, latent_image, latent.get("downscale_ratio_spacial", None))
latent_image = comfy.sample.fix_empty_latent_channels(model, latent_image, latent.get("downscale_ratio_spacial", None), latent.get("downscale_ratio_temporal", None))
latent["samples"] = latent_image
if not add_noise:
@ -770,6 +770,7 @@ class SamplerCustom(io.ComfyNode):
out = latent.copy()
out.pop("downscale_ratio_spacial", None)
out.pop("downscale_ratio_temporal", None)
out["samples"] = samples
if "x0" in x0_output:
x0_out = model.model.process_latent_out(x0_output["x0"].cpu())
@ -793,7 +794,8 @@ class BasicGuider(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="BasicGuider",
category="sampling/custom_sampling/guiders",
display_name="Basic Guider",
category="sampling/guiders",
inputs=[
io.Model.Input("model"),
io.Conditioning.Input("conditioning"),
@ -814,7 +816,8 @@ class CFGGuider(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="CFGGuider",
category="sampling/custom_sampling/guiders",
display_name="CFG Guider",
category="sampling/guiders",
inputs=[
io.Model.Input("model"),
io.Conditioning.Input("positive"),
@ -868,7 +871,8 @@ class DualCFGGuider(io.ComfyNode):
return io.Schema(
node_id="DualCFGGuider",
search_aliases=["dual prompt guidance"],
category="sampling/custom_sampling/guiders",
display_name="Dual CFG Guider",
category="sampling/guiders",
inputs=[
io.Model.Input("model"),
io.Conditioning.Input("cond1"),
@ -896,7 +900,7 @@ class DisableNoise(io.ComfyNode):
return io.Schema(
node_id="DisableNoise",
search_aliases=["zero noise"],
category="sampling/custom_sampling/noise",
category="sampling/noise",
inputs=[],
outputs=[io.Noise.Output()]
)
@ -913,7 +917,7 @@ class RandomNoise(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="RandomNoise",
category="sampling/custom_sampling/noise",
category="sampling/noise",
inputs=[io.Int.Input("noise_seed", default=0, min=0, max=0xffffffffffffffff, control_after_generate=True)],
outputs=[io.Noise.Output()]
)
@ -949,7 +953,7 @@ class SamplerCustomAdvanced(io.ComfyNode):
latent = latent_image
latent_image = latent["samples"]
latent = latent.copy()
latent_image = comfy.sample.fix_empty_latent_channels(guider.model_patcher, latent_image, latent.get("downscale_ratio_spacial", None))
latent_image = comfy.sample.fix_empty_latent_channels(guider.model_patcher, latent_image, latent.get("downscale_ratio_spacial", None), latent.get("downscale_ratio_temporal", None))
latent["samples"] = latent_image
noise_mask = None
@ -965,6 +969,7 @@ class SamplerCustomAdvanced(io.ComfyNode):
out = latent.copy()
out.pop("downscale_ratio_spacial", None)
out.pop("downscale_ratio_temporal", None)
out["samples"] = samples
if "x0" in x0_output:
x0_out = guider.model_patcher.model.process_latent_out(x0_output["x0"].cpu())

View File

@ -215,7 +215,7 @@ class Flux2Scheduler(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="Flux2Scheduler",
category="sampling/custom_sampling/schedulers",
category="sampling/schedulers",
inputs=[
io.Int.Input("steps", default=20, min=1, max=4096),
io.Int.Input("width", default=1024, min=16, max=nodes.MAX_RESOLUTION, step=1),
@ -263,7 +263,7 @@ class FluxKVCache(io.ComfyNode):
node_id="FluxKVCache",
display_name="Flux KV Cache",
description="Enables KV Cache optimization for reference images on Flux family models.",
category="",
category="experimental",
is_experimental=True,
inputs=[
io.Model.Input("model", tooltip="The model to use KV Cache on."),

View File

@ -340,7 +340,7 @@ class GITSScheduler(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="GITSScheduler",
category="sampling/custom_sampling/schedulers",
category="sampling/schedulers",
inputs=[
io.Float.Input("coeff", default=1.20, min=0.80, max=1.50, step=0.05, advanced=True),
io.Int.Input("steps", default=10, min=2, max=1000),

View File

@ -162,7 +162,7 @@ class ImageAddNoise(IO.ComfyNode):
node_id="ImageAddNoise",
search_aliases=["film grain"],
display_name="Add Noise to Image",
category="image/postprocessing",
category="image/filters",
inputs=[
IO.Image.Input("image"),
IO.Int.Input(
@ -194,7 +194,8 @@ class SaveAnimatedWEBP(IO.ComfyNode):
def define_schema(cls):
return IO.Schema(
node_id="SaveAnimatedWEBP",
category="image/animation",
display_name="Save Animated WEBP",
category="image",
inputs=[
IO.Image.Input("images"),
IO.String.Input("filename_prefix", default="ComfyUI"),
@ -231,7 +232,8 @@ class SaveAnimatedPNG(IO.ComfyNode):
def define_schema(cls):
return IO.Schema(
node_id="SaveAnimatedPNG",
category="image/animation",
display_name="Save Animated PNG",
category="image",
inputs=[
IO.Image.Input("images"),
IO.String.Input("filename_prefix", default="ComfyUI"),
@ -493,7 +495,7 @@ class SaveSVGNode(IO.ComfyNode):
search_aliases=["export vector", "save vector graphics"],
display_name="Save SVG",
description="Save SVG files on disk.",
category="image/save",
category="image",
inputs=[
IO.SVG.Input("svg"),
IO.String.Input(

View File

@ -175,7 +175,7 @@ class LTXVImgToVideoInplace(io.ComfyNode):
generate = execute # TODO: remove
def _append_guide_attention_entry(positive, negative, pre_filter_count, latent_shape, strength=1.0):
def _append_guide_attention_entry(positive, negative, pre_filter_count, latent_shape, strength=1.0, attention_mask=None):
"""Append a guide_attention_entry to both positive and negative conditioning.
Each entry tracks one guide reference for per-reference attention control.
@ -184,9 +184,10 @@ def _append_guide_attention_entry(positive, negative, pre_filter_count, latent_s
new_entry = {
"pre_filter_count": pre_filter_count,
"strength": strength,
"pixel_mask": None,
"pixel_mask": attention_mask.unsqueeze(0).unsqueeze(0) if attention_mask is not None else None, # reshape to (1, 1, F, H, W)
"latent_shape": latent_shape,
}
results = []
for cond in (positive, negative):
# Read existing entries from this specific conditioning
@ -196,8 +197,7 @@ def _append_guide_attention_entry(positive, negative, pre_filter_count, latent_s
if found is not None:
existing = found
break
# Shallow copy and append (no deepcopy needed — entries contain
# only scalars and None for pixel_mask at this call site).
# Shallow copy only and append (pixel_mask is never mutated).
entries = [*existing, new_entry]
results.append(node_helpers.conditioning_set_values(
cond, {"guide_attention_entries": entries}
@ -263,6 +263,12 @@ class LTXVAddGuide(io.ComfyNode):
"down to the nearest multiple of 8. Negative values are counted from the end of the video.",
),
io.Float.Input("strength", default=1.0, min=0.0, max=10.0, step=0.01),
io.Mask.Input(
"attention_mask",
optional=True,
tooltip="Optional pixel-space spatial mask. Controls per-region "
"conditioning influence via self-attention, multiplied by strength.",
),
ICLoRAParameters.Input(
"iclora_parameters",
optional=True,
@ -410,7 +416,7 @@ class LTXVAddGuide(io.ComfyNode):
return latent_image, noise_mask
@classmethod
def execute(cls, positive, negative, vae, latent, image, frame_idx, strength, iclora_parameters=None) -> io.NodeOutput:
def execute(cls, positive, negative, vae, latent, image, frame_idx, strength, attention_mask=None, iclora_parameters=None) -> io.NodeOutput:
scale_factors = vae.downscale_index_formula
latent_image = latent["samples"]
noise_mask = get_noise_mask(latent)
@ -469,6 +475,7 @@ class LTXVAddGuide(io.ComfyNode):
pre_filter_count = t.shape[2] * t.shape[3] * t.shape[4]
positive, negative = _append_guide_attention_entry(
positive, negative, pre_filter_count, guide_latent_shape, strength=strength,
attention_mask=attention_mask,
)
return io.NodeOutput(positive, negative, {"samples": latent_image, "noise_mask": noise_mask})
@ -594,7 +601,7 @@ class LTXVScheduler(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="LTXVScheduler",
category="sampling/custom_sampling/schedulers",
category="sampling/schedulers",
inputs=[
io.Int.Input("steps", default=20, min=1, max=10000),
io.Float.Input("max_shift", default=2.05, min=0.0, max=100.0, step=0.01),

View File

@ -83,7 +83,7 @@ class ImageCompositeMasked(IO.ComfyNode):
node_id="ImageCompositeMasked",
search_aliases=["overlay", "layer", "paste image", "images composition"],
display_name="Image Composite Masked",
category="image",
category="image/compositing",
inputs=[
IO.Image.Input("destination"),
IO.Image.Input("source"),
@ -112,7 +112,7 @@ class MaskToImage(IO.ComfyNode):
node_id="MaskToImage",
search_aliases=["convert mask"],
display_name="Convert Mask to Image",
category="mask",
category="image/mask",
inputs=[
IO.Mask.Input("mask"),
],
@ -134,7 +134,7 @@ class ImageToMask(IO.ComfyNode):
node_id="ImageToMask",
search_aliases=["extract channel", "channel to mask"],
display_name="Convert Image to Mask",
category="mask",
category="image/mask",
inputs=[
IO.Image.Input("image"),
IO.Combo.Input("channel", options=["red", "green", "blue", "alpha"]),
@ -157,7 +157,8 @@ class ImageColorToMask(IO.ComfyNode):
return IO.Schema(
node_id="ImageColorToMask",
search_aliases=["color keying", "chroma key"],
category="mask",
display_name="Convert Image Color to Mask",
category="image/mask",
inputs=[
IO.Image.Input("image"),
IO.Int.Input("color", default=0, min=0, max=0xFFFFFF, step=1, display_mode=IO.NumberDisplay.number),
@ -180,7 +181,8 @@ class SolidMask(IO.ComfyNode):
def define_schema(cls):
return IO.Schema(
node_id="SolidMask",
category="mask",
display_name="Create Solid Mask",
category="image/mask",
inputs=[
IO.Float.Input("value", default=1.0, min=0.0, max=1.0, step=0.01),
IO.Int.Input("width", default=512, min=1, max=nodes.MAX_RESOLUTION, step=1),
@ -204,7 +206,7 @@ class InvertMask(IO.ComfyNode):
node_id="InvertMask",
search_aliases=["reverse mask", "flip mask"],
display_name="Invert Mask",
category="mask",
category="image/mask",
inputs=[
IO.Mask.Input("mask"),
],
@ -226,7 +228,7 @@ class CropMask(IO.ComfyNode):
node_id="CropMask",
search_aliases=["cut mask", "extract mask region", "mask slice"],
display_name="Crop Mask",
category="mask",
category="image/mask",
inputs=[
IO.Mask.Input("mask"),
IO.Int.Input("x", default=0, min=0, max=nodes.MAX_RESOLUTION, step=1),
@ -253,7 +255,7 @@ class MaskComposite(IO.ComfyNode):
node_id="MaskComposite",
search_aliases=["combine masks", "blend masks", "layer masks", "masks composition"],
display_name="Combine Masks",
category="mask",
category="image/mask",
inputs=[
IO.Mask.Input("destination"),
IO.Mask.Input("source"),
@ -304,7 +306,7 @@ class FeatherMask(IO.ComfyNode):
node_id="FeatherMask",
search_aliases=["soft edge mask", "blur mask edges", "gradient mask edge"],
display_name="Feather Mask",
category="mask",
category="image/mask",
inputs=[
IO.Mask.Input("mask"),
IO.Int.Input("left", default=0, min=0, max=nodes.MAX_RESOLUTION, step=1),
@ -352,7 +354,7 @@ class GrowMask(IO.ComfyNode):
node_id="GrowMask",
search_aliases=["expand mask", "shrink mask"],
display_name="Grow Mask",
category="mask",
category="image/mask",
inputs=[
IO.Mask.Input("mask"),
IO.Int.Input("expand", default=0, min=-nodes.MAX_RESOLUTION, max=nodes.MAX_RESOLUTION, step=1),
@ -388,7 +390,8 @@ class ThresholdMask(IO.ComfyNode):
return IO.Schema(
node_id="ThresholdMask",
search_aliases=["binary mask"],
category="mask",
display_name="Threshold Mask",
category="image/mask",
inputs=[
IO.Mask.Input("mask"),
IO.Float.Input("value", default=0.5, min=0.0, max=1.0, step=0.01),
@ -414,7 +417,7 @@ class MaskPreview(IO.ComfyNode):
node_id="MaskPreview",
search_aliases=["show mask", "view mask", "inspect mask", "debug mask"],
display_name="Preview Mask",
category="mask",
category="image/mask",
description="Saves the input images to your ComfyUI output directory.",
inputs=[
IO.Mask.Input("mask"),

View File

@ -276,8 +276,8 @@ class CLIPSave:
for x in extra_pnginfo:
metadata[x] = json.dumps(extra_pnginfo[x])
comfy.model_management.load_models_gpu([clip.load_model()], force_patch_weights=True)
clip_sd = clip.get_sd()
clip.load_model()
clip_sd = clip.state_dict_for_saving()
for prefix in ["clip_l.", "clip_g.", "clip_h.", "t5xxl.", "pile_t5xl.", "mt5xl.", "umt5xxl.", "t5base.", "gemma2_2b.", "llama.", "hydit_clip.", ""]:
k = list(filter(lambda a: a.startswith(prefix), clip_sd.keys()))

View File

@ -13,8 +13,8 @@ class Morphology(io.ComfyNode):
return io.Schema(
node_id="Morphology",
search_aliases=["erode", "dilate"],
display_name="ImageMorphology",
category="image/postprocessing",
display_name="Apply Morphology",
category="image/filters",
inputs=[
io.Image.Input("image"),
io.Combo.Input(

View File

@ -13,7 +13,7 @@ class wanBlockSwap(io.ComfyNode):
return io.Schema(
node_id="wanBlockSwap",
category="",
description="NOP",
description="Intercept wanBlockSwap custom node that causes major instability and make it no-op.",
inputs=[
io.Model.Input("model"),
],

View File

@ -20,7 +20,7 @@ class NumberConvertNode(io.ComfyNode):
def define_schema(cls) -> io.Schema:
return io.Schema(
node_id="ComfyNumberConvert",
display_name="Number Convert",
display_name="Convert Number",
category="utils",
search_aliases=[
"int to float", "float to int", "number convert",

View File

@ -31,7 +31,7 @@ class OptimalStepsScheduler(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="OptimalStepsScheduler",
category="sampling/custom_sampling/schedulers",
category="sampling/schedulers",
inputs=[
io.Combo.Input("model_type", options=["FLUX", "Wan", "Chroma"]),
io.Int.Input("steps", default=20, min=3, max=1000),

View File

@ -22,7 +22,7 @@ class Blend(io.ComfyNode):
node_id="ImageBlend",
search_aliases=["mix images"],
display_name="Blend Images",
category="image/postprocessing",
category="image/filters",
essentials_category="Image Tools",
inputs=[
io.Image.Input("image1"),
@ -80,8 +80,8 @@ class Blur(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="ImageBlur",
display_name="Image Blur",
category="image/postprocessing",
display_name="Blur Image",
category="image/filters",
inputs=[
io.Image.Input("image"),
io.Int.Input("blur_radius", default=1, min=1, max=31, step=1),
@ -117,7 +117,7 @@ class Quantize(io.ComfyNode):
return io.Schema(
node_id="ImageQuantize",
display_name="Quantize Image",
category="image/postprocessing",
category="image/filters",
inputs=[
io.Image.Input("image"),
io.Int.Input("colors", default=256, min=1, max=256, step=1),
@ -183,7 +183,7 @@ class Sharpen(io.ComfyNode):
return io.Schema(
node_id="ImageSharpen",
display_name="Sharpen Image",
category="image/postprocessing",
category="image/filters",
inputs=[
io.Image.Input("image"),
io.Int.Input("sharpen_radius", default=1, min=1, max=31, step=1, advanced=True),
@ -568,7 +568,7 @@ def batch_latents(latents: list[dict[str, torch.Tensor]]) -> dict[str, torch.Ten
class BatchImagesNode(io.ComfyNode):
@classmethod
def define_schema(cls):
autogrow_template = io.Autogrow.TemplatePrefix(io.Image.Input("image"), prefix="image", min=2, max=50)
autogrow_template = io.Autogrow.TemplatePrefix(io.Image.Input("image"), prefix="image", min=1, max=50)
return io.Schema(
node_id="BatchImagesNode",
display_name="Batch Images",
@ -590,12 +590,12 @@ class BatchImagesNode(io.ComfyNode):
class BatchMasksNode(io.ComfyNode):
@classmethod
def define_schema(cls):
autogrow_template = io.Autogrow.TemplatePrefix(io.Mask.Input("mask"), prefix="mask", min=2, max=50)
autogrow_template = io.Autogrow.TemplatePrefix(io.Mask.Input("mask"), prefix="mask", min=1, max=50)
return io.Schema(
node_id="BatchMasksNode",
search_aliases=["combine masks", "stack masks", "merge masks"],
display_name="Batch Masks",
category="mask",
category="image/mask",
inputs=[
io.Autogrow.Input("masks", template=autogrow_template)
],
@ -611,7 +611,7 @@ class BatchMasksNode(io.ComfyNode):
class BatchLatentsNode(io.ComfyNode):
@classmethod
def define_schema(cls):
autogrow_template = io.Autogrow.TemplatePrefix(io.Latent.Input("latent"), prefix="latent", min=2, max=50)
autogrow_template = io.Autogrow.TemplatePrefix(io.Latent.Input("latent"), prefix="latent", min=1, max=50)
return io.Schema(
node_id="BatchLatentsNode",
search_aliases=["combine latents", "stack latents", "merge latents"],
@ -670,8 +670,8 @@ class ColorTransfer(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="ColorTransfer",
display_name="Color Transfer",
category="image/postprocessing",
display_name="Transfer Color",
category="image/filters",
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=[

View File

@ -15,7 +15,7 @@ class RTDETR_detect(io.ComfyNode):
return io.Schema(
node_id="RTDETR_detect",
display_name="RT-DETR Detect",
category="detection",
category="image/detection",
search_aliases=["bbox", "bounding box", "object detection", "coco"],
inputs=[
io.Model.Input("model", display_name="model"),
@ -71,7 +71,7 @@ class DrawBBoxes(io.ComfyNode):
return io.Schema(
node_id="DrawBBoxes",
display_name="Draw BBoxes",
category="detection",
category="image/detection",
search_aliases=["bbox", "bounding box", "object detection", "rt_detr", "visualize detections", "coco"],
inputs=[
io.Image.Input("image", optional=True),

View File

@ -93,7 +93,7 @@ class SAM3_Detect(io.ComfyNode):
return io.Schema(
node_id="SAM3_Detect",
display_name="SAM3 Detect",
category="detection",
category="image/detection",
search_aliases=["sam3", "segment anything", "open vocabulary", "text detection", "segment"],
inputs=[
io.Model.Input("model", display_name="model"),
@ -265,7 +265,7 @@ class SAM3_VideoTrack(io.ComfyNode):
return io.Schema(
node_id="SAM3_VideoTrack",
display_name="SAM3 Video Track",
category="detection",
category="image/detection",
search_aliases=["sam3", "video", "track", "propagate"],
inputs=[
io.Image.Input("images", display_name="images", tooltip="Video frames as batched images"),
@ -320,7 +320,7 @@ class SAM3_TrackPreview(io.ComfyNode):
return io.Schema(
node_id="SAM3_TrackPreview",
display_name="SAM3 Track Preview",
category="detection",
category="image/detection",
inputs=[
SAM3TrackData.Input("track_data", display_name="track_data"),
io.Image.Input("images", display_name="images", optional=True),
@ -478,7 +478,7 @@ class SAM3_TrackToMask(io.ComfyNode):
return io.Schema(
node_id="SAM3_TrackToMask",
display_name="SAM3 Track to Mask",
category="detection",
category="image/detection",
inputs=[
SAM3TrackData.Input("track_data", display_name="track_data"),
io.String.Input("object_indices", display_name="object_indices", default="",

View File

@ -353,7 +353,8 @@ class SDPoseDrawKeypoints(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="SDPoseDrawKeypoints",
category="image/preprocessors",
display_name="SDPose Draw Keypoints",
category="image/detection",
search_aliases=["openpose", "pose detection", "preprocessor", "keypoints", "pose"],
inputs=[
io.Custom("POSE_KEYPOINT").Input("keypoints"),
@ -421,7 +422,8 @@ class SDPoseKeypointExtractor(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="SDPoseKeypointExtractor",
category="image/preprocessors",
display_name="SDPose Keypoint Extractor",
category="image/detection",
search_aliases=["openpose", "pose detection", "preprocessor", "keypoints", "sdpose"],
description="Extract pose keypoints from images using the SDPose model: https://huggingface.co/Comfy-Org/SDPose/tree/main/checkpoints",
inputs=[
@ -595,7 +597,8 @@ class SDPoseFaceBBoxes(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="SDPoseFaceBBoxes",
category="image/preprocessors",
display_name="SDPose Face Bounding Boxes",
category="image/detection",
search_aliases=["face bbox", "face bounding box", "pose", "keypoints"],
inputs=[
io.Custom("POSE_KEYPOINT").Input("keypoints"),
@ -652,7 +655,8 @@ class CropByBBoxes(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="CropByBBoxes",
category="image/preprocessors",
display_name="Crop By Bounding Boxes",
category="image/transform",
search_aliases=["crop", "face crop", "bbox crop", "pose", "bounding box"],
description="Crop and resize regions from the input image batch based on provided bounding boxes.",
inputs=[

View File

@ -65,7 +65,7 @@ class VideoLinearCFGGuidance:
RETURN_TYPES = ("MODEL",)
FUNCTION = "patch"
CATEGORY = "sampling/video_models"
CATEGORY = "sampling/guiders"
def patch(self, model, min_cfg):
def linear_cfg(args):
@ -89,7 +89,7 @@ class VideoTriangleCFGGuidance:
RETURN_TYPES = ("MODEL",)
FUNCTION = "patch"
CATEGORY = "sampling/video_models"
CATEGORY = "sampling/guiders"
def patch(self, model, min_cfg):
def linear_cfg(args):
@ -157,5 +157,7 @@ NODE_CLASS_MAPPINGS = {
}
NODE_DISPLAY_NAME_MAPPINGS = {
"ImageOnlyCheckpointLoader": "Image Only Checkpoint Loader (img2vid model)",
"ImageOnlyCheckpointLoader": "Load Checkpoint Image Only (img2vid model)",
"VideoLinearCFGGuidance": "Video Linear CFG Guidance",
"VideoTriangleCFGGuidance": "Video Triangle CFG Guidance",
}

View File

@ -122,7 +122,8 @@ class VOIDQuadmaskPreprocess(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="VOIDQuadmaskPreprocess",
category="mask/video",
display_name="VOID Quadmask Preprocessor",
category="image/mask",
inputs=[
io.Mask.Input("mask"),
io.Int.Input("dilate_width", default=0, min=0, max=50, step=1,
@ -392,7 +393,7 @@ class VOIDWarpedNoiseSource(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="VOIDWarpedNoiseSource",
category="sampling/custom_sampling/noise",
category="sampling/noise",
inputs=[
io.Latent.Input("warped_noise",
tooltip="Warped noise latent from VOIDWarpedNoise"),
@ -454,7 +455,7 @@ class VOIDSampler(io.ComfyNode):
def define_schema(cls):
return io.Schema(
node_id="VOIDSampler",
category="sampling/custom_sampling/samplers",
category="sampling/samplers",
inputs=[],
outputs=[io.Sampler.Output()],
)

View File

@ -691,7 +691,7 @@ class LoraLoader:
FUNCTION = "load_lora"
CATEGORY = "loaders"
DESCRIPTION = "LoRAs are used to modify diffusion and CLIP models, altering the way in which latents are denoised such as applying styles. Multiple LoRA nodes can be linked together."
DESCRIPTION = "This LoRA loader is used to modify both diffusion and CLIP models, altering the way in which latents are denoised such as applying styles. Multiple LoRA nodes can be linked together."
SEARCH_ALIASES = ["lora", "load lora", "apply lora", "lora loader", "lora model"]
def load_lora(self, model, clip, lora_name, strength_model, strength_clip):
@ -723,6 +723,7 @@ class LoraLoaderModelOnly(LoraLoader):
"strength_model": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}),
}}
RETURN_TYPES = ("MODEL",)
DESCRIPTION = "This LoRAs loader is used to modify the diffusion model, altering the way in which latents are denoised such as applying styles. Multiple LoRA nodes can be linked together."
FUNCTION = "load_lora_model_only"
def load_lora_model_only(self, model, lora_name, strength_model):
@ -1524,7 +1525,7 @@ class SetLatentNoiseMask:
def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, denoise=1.0, disable_noise=False, start_step=None, last_step=None, force_full_denoise=False):
latent_image = latent["samples"]
latent_image = comfy.sample.fix_empty_latent_channels(model, latent_image, latent.get("downscale_ratio_spacial", None))
latent_image = comfy.sample.fix_empty_latent_channels(model, latent_image, latent.get("downscale_ratio_spacial", None), latent.get("downscale_ratio_temporal", None))
if disable_noise:
noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
@ -1543,6 +1544,7 @@ def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive,
force_full_denoise=force_full_denoise, noise_mask=noise_mask, callback=callback, disable_pbar=disable_pbar, seed=seed)
out = latent.copy()
out.pop("downscale_ratio_spacial", None)
out.pop("downscale_ratio_temporal", None)
out["samples"] = samples
return (out, )
@ -1775,7 +1777,7 @@ class LoadImageMask(LoadImage):
}
}
CATEGORY = "mask"
CATEGORY = "image"
RETURN_TYPES = ("MASK",)
FUNCTION = "load_image_mask"

View File

@ -485,8 +485,15 @@ paths:
post:
operationId: uploadMask
tags: [upload]
summary: Upload a mask image
description: Uploads a mask image associated with a previously-uploaded reference image.
deprecated: true
summary: Upload a mask image (deprecated)
description: |
Deprecated. Clients should composite the mask onto the source image
client-side and upload the resulting image via POST /api/upload/image
instead. This endpoint will continue to function for older clients,
but will not receive new features.
Uploads a mask image associated with a previously-uploaded reference image.
requestBody:
required: true
content:

View File

@ -8,11 +8,44 @@ from app.model_download import (
ModelDownloadRequest,
is_allowed_model_download_url,
normalize_model_relative_path,
open_model_download_response,
parse_model_download_request,
resolve_model_download_destination,
)
class _FakeResponse:
"""Minimal stand-in for ``aiohttp.ClientResponse`` for the redirect tests."""
def __init__(self, status, headers=None):
self.status = status
self.headers = headers or {}
self.released = False
def release(self):
self.released = True
async def __aenter__(self):
return self
async def __aexit__(self, exc_type, exc, tb):
self.released = True
class _FakeSession:
"""Hands out queued ``_FakeResponse`` objects in order."""
def __init__(self, responses):
self._responses = list(responses)
self.calls = []
async def get(self, url, allow_redirects, timeout):
self.calls.append((url, allow_redirects))
if not self._responses:
raise AssertionError("Unexpected extra session.get call")
return self._responses.pop(0)
def test_parse_model_download_request_allows_huggingface_model_url():
request = parse_model_download_request({
"name": "nested/model.safetensors",
@ -33,12 +66,46 @@ def test_parse_model_download_request_allows_huggingface_model_url():
"http://localhost:8000/model.safetensors",
"http://huggingface.co/org/repo/resolve/main/model.safetensors",
"https://example.com/model.safetensors",
"https://huggingface.co.evil.com/model.safetensors",
],
)
def test_download_url_allowlist_rejects_untrusted_or_plain_http_urls(url):
assert is_allowed_model_download_url(url) is False
@pytest.mark.parametrize(
"url",
[
# Direct HF model URLs.
"https://huggingface.co/org/repo/resolve/main/model.safetensors",
# HF LFS CDN subdomains: this is where `/resolve/main/...` redirects
# land, so the allowlist must accept them or downloads break.
"https://cdn-lfs.huggingface.co/repos/abc/def/model.safetensors",
"https://cdn-lfs-us-1.huggingface.co/repos/abc/def/model.safetensors",
# Civitai download endpoints (PR objective: support Civitai too).
"https://civitai.com/api/download/models/12345",
"https://civitai.red/api/download/models/12345",
],
)
def test_download_url_allowlist_accepts_huggingface_and_civitai_urls(url):
assert is_allowed_model_download_url(url) is True
@pytest.mark.parametrize(
"name, expected",
[
("model.safetensors", "model.safetensors"),
("sub/model.safetensors", "sub/model.safetensors"),
("nested/dir/model.safetensors", "nested/dir/model.safetensors"),
# Backslashes are normalized to forward slashes so Windows-style
# paths land in the same place as the POSIX equivalents.
("nested\\dir\\model.safetensors", "nested/dir/model.safetensors"),
],
)
def test_normalize_model_relative_path_accepts_safe_paths(name, expected):
assert normalize_model_relative_path(name) == expected
@pytest.mark.parametrize(
"name",
[
@ -112,3 +179,66 @@ def test_resolve_model_download_destination_rejects_blocked_or_unknown_directori
url="https://huggingface.co/org/repo/resolve/main/model.safetensors",
directory=directory,
))
@pytest.mark.asyncio
async def test_open_model_download_response_follows_allowed_subdomain_redirect():
"""HF redirects /resolve/main/... to cdn-lfs.huggingface.co; that must work."""
session = _FakeSession([
_FakeResponse(302, {"Location": "https://cdn-lfs.huggingface.co/repos/abc/model.safetensors"}),
_FakeResponse(200),
])
response = await open_model_download_response(
session, "https://huggingface.co/org/repo/resolve/main/model.safetensors"
)
assert response.status == 200
assert session.calls == [
("https://huggingface.co/org/repo/resolve/main/model.safetensors", False),
("https://cdn-lfs.huggingface.co/repos/abc/model.safetensors", False),
]
@pytest.mark.asyncio
async def test_open_model_download_response_rejects_offsite_redirect():
"""A redirect leaving the allowlist must surface as a 403 instead of being followed."""
session = _FakeSession([
_FakeResponse(302, {"Location": "https://attacker.example.com/payload"}),
])
with pytest.raises(ModelDownloadError) as exc_info:
await open_model_download_response(
session, "https://huggingface.co/org/repo/resolve/main/model.safetensors"
)
assert exc_info.value.status == 403
# The initial request was issued with redirects disabled, otherwise
# the validation above would be a no-op.
assert session.calls[0][1] is False
@pytest.mark.asyncio
async def test_open_model_download_response_rejects_redirect_without_location():
session = _FakeSession([_FakeResponse(302)])
with pytest.raises(ModelDownloadError) as exc_info:
await open_model_download_response(
session, "https://huggingface.co/org/repo/resolve/main/model.safetensors"
)
assert exc_info.value.status == 502
@pytest.mark.asyncio
async def test_open_model_download_response_stops_after_too_many_redirects():
session = _FakeSession(
[_FakeResponse(302, {"Location": "https://cdn-lfs.huggingface.co/loop"})] * 10
)
with pytest.raises(ModelDownloadError) as exc_info:
await open_model_download_response(
session, "https://huggingface.co/org/repo/resolve/main/model.safetensors"
)
assert exc_info.value.status == 502