mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-21 23:41:28 +08:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0e14070ee3 | ||
|
|
65eac2ba82 | ||
|
|
08a4be9508 |
@@ -32,11 +32,9 @@ jobs:
|
||||
PR_NUMBER: ${{ github.event.pull_request.number || github.event.issue.number }}
|
||||
PR_AUTHOR: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
|
||||
BASE_ALLOWLIST: action@github.com,actions-user,ampagent,claude,comfy-pr-bot,GitHub Action,github-actions,github-actions[bot],Glary Bot,Glary-Bot,*[bot]
|
||||
# For each commit emit the GitHub login when the author/committer email resolves to a GitHub account
|
||||
# otherwise fall back to the raw git name.
|
||||
run: |
|
||||
others=$(gh api "repos/${{ github.repository }}/pulls/${PR_NUMBER}/commits" --paginate \
|
||||
--jq '.[] | (.author.login // .commit.author.name // empty), (.committer.login // .commit.committer.name // empty)' \
|
||||
--jq '.[] | (.author.login // empty), (.committer.login // empty)' \
|
||||
| sort -u | grep -vix "${PR_AUTHOR}" | paste -sd, -)
|
||||
if [ -n "$others" ]; then
|
||||
echo "allowlist=${BASE_ALLOWLIST},${others}" >> "$GITHUB_OUTPUT"
|
||||
@@ -45,7 +43,7 @@ jobs:
|
||||
fi
|
||||
|
||||
- name: CLA Assistant
|
||||
# Run on PR events, on "recheck" comment, or when someone posts the signing phrase.
|
||||
# Run on PR events, on "recheck" comment, or when someone posts the exact signing phrase.
|
||||
# IMPORTANT: this phrase must match `custom-pr-sign-comment` below.
|
||||
if: >
|
||||
github.event_name == 'pull_request_target' ||
|
||||
|
||||
@@ -228,6 +228,7 @@ async def list_assets_route(request: web.Request) -> web.Response:
|
||||
exclude_tags=q.exclude_tags,
|
||||
name_contains=q.name_contains,
|
||||
metadata_filter=q.metadata_filter,
|
||||
asset_hash=q.hash,
|
||||
limit=q.limit,
|
||||
offset=q.offset,
|
||||
sort=sort,
|
||||
|
||||
@@ -54,6 +54,18 @@ class ListAssetsQuery(BaseModel):
|
||||
exclude_tags: list[str] = Field(default_factory=list)
|
||||
name_contains: str | None = None
|
||||
|
||||
# Filter to assets whose content hash matches exactly. Param name is `hash`
|
||||
# per the projected openapi.yaml listAssets contract (the response-body field
|
||||
# is `asset_hash`; the query param is `hash`).
|
||||
hash: str | None = None
|
||||
|
||||
# Declared for cloud/core contract parity. In core, reads are owner-scoped
|
||||
# (owner_id == "") and there is no separate shared/public pool for this flag
|
||||
# to include or exclude, so it is inert here and intentionally not threaded
|
||||
# into the query. Accepted (not rejected) so the FE needs no isCloud branch;
|
||||
# cloud enforces the flag in its own service layer.
|
||||
include_public: bool = True
|
||||
|
||||
# Accept either a JSON string (query param) or a dict
|
||||
metadata_filter: dict[str, Any] | None = None
|
||||
|
||||
@@ -86,6 +98,19 @@ class ListAssetsQuery(BaseModel):
|
||||
return out
|
||||
return v
|
||||
|
||||
@field_validator("hash", mode="before")
|
||||
@classmethod
|
||||
def _normalize_hash(cls, v):
|
||||
# Normalize for an exact match against stored hashes (which are
|
||||
# lowercase `blake3:<hex>`). Liberal in what we accept — no pattern
|
||||
# enforcement; a non-matching value simply yields an empty page.
|
||||
# An explicitly-supplied-but-empty value (`?hash=`) stays `""` so it
|
||||
# is treated as an exact-match miss (empty page), not silently dropped
|
||||
# to "no filter" — omit the param entirely to disable the filter.
|
||||
if isinstance(v, str):
|
||||
return v.strip().lower()
|
||||
return v
|
||||
|
||||
@field_validator("metadata_filter", mode="before")
|
||||
@classmethod
|
||||
def _parse_metadata_json(cls, v):
|
||||
|
||||
@@ -261,6 +261,7 @@ def list_references_page(
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
name_contains: str | None = None,
|
||||
asset_hash: str | None = None,
|
||||
include_tags: Sequence[str] | None = None,
|
||||
exclude_tags: Sequence[str] | None = None,
|
||||
metadata_filter: dict | None = None,
|
||||
@@ -293,6 +294,11 @@ def list_references_page(
|
||||
escaped, esc = escape_sql_like_string(name_contains)
|
||||
base = base.where(AssetReference.name.ilike(f"%{escaped}%", escape=esc))
|
||||
|
||||
# `is not None` (not truthiness): an explicit empty hash is an exact-match
|
||||
# miss (empty page), while an omitted hash (None) disables the filter.
|
||||
if asset_hash is not None:
|
||||
base = base.where(Asset.hash == asset_hash)
|
||||
|
||||
base = apply_tag_filters(base, include_tags, exclude_tags)
|
||||
base = apply_metadata_filter(base, metadata_filter)
|
||||
|
||||
@@ -345,6 +351,8 @@ def list_references_page(
|
||||
count_stmt = count_stmt.where(
|
||||
AssetReference.name.ilike(f"%{escaped}%", escape=esc)
|
||||
)
|
||||
if asset_hash is not None:
|
||||
count_stmt = count_stmt.where(Asset.hash == asset_hash)
|
||||
count_stmt = apply_tag_filters(count_stmt, include_tags, exclude_tags)
|
||||
count_stmt = apply_metadata_filter(count_stmt, metadata_filter)
|
||||
|
||||
|
||||
@@ -274,6 +274,7 @@ def list_assets_page(
|
||||
exclude_tags: Sequence[str] | None = None,
|
||||
name_contains: str | None = None,
|
||||
metadata_filter: dict | None = None,
|
||||
asset_hash: str | None = None,
|
||||
limit: int = 20,
|
||||
offset: int = 0,
|
||||
sort: str = "created_at",
|
||||
@@ -319,6 +320,7 @@ def list_assets_page(
|
||||
exclude_tags=exclude_tags,
|
||||
name_contains=name_contains,
|
||||
metadata_filter=metadata_filter,
|
||||
asset_hash=asset_hash,
|
||||
limit=fetch_limit,
|
||||
offset=offset,
|
||||
sort=sort,
|
||||
|
||||
@@ -35,11 +35,7 @@ class ModelFileManager:
|
||||
for folder in model_types:
|
||||
if folder in folder_black_list:
|
||||
continue
|
||||
output_folders.append({
|
||||
"name": folder,
|
||||
"folders": folder_paths.get_folder_paths(folder),
|
||||
"extensions": sorted(folder_paths.folder_names_and_paths[folder][1]),
|
||||
})
|
||||
output_folders.append({"name": folder, "folders": folder_paths.get_folder_paths(folder)})
|
||||
return web.json_response(output_folders)
|
||||
|
||||
# NOTE: This is an experiment to replace `/models/{folder}`
|
||||
|
||||
@@ -92,7 +92,6 @@ 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.")
|
||||
parser.add_argument("--disable-triton-backend", action="store_true", help="Force-disable the comfy-kitchen Triton backend, overriding the automatic ROCm/AMD default and --enable-triton-backend.")
|
||||
|
||||
class LatentPreviewMethod(enum.Enum):
|
||||
NoPreviews = "none"
|
||||
|
||||
@@ -1,46 +0,0 @@
|
||||
"""Runtime config the frontend reads from /features to follow --comfy-api-base.
|
||||
|
||||
For a non-prod comfy.org backend (staging or an ephemeral preview env), "/features" exposes the api and
|
||||
platform base so the frontend talks to it without a rebuild, plus the Firebase environment it should use.
|
||||
Prod bases are left alone and keep their build-time defaults.
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
from urllib.parse import urlparse
|
||||
|
||||
from comfy.cli_args import args
|
||||
|
||||
_STAGING_API_HOST = "stagingapi.comfy.org"
|
||||
_TESTENV_HOST_SUFFIX = ".testenvs.comfy.org"
|
||||
_STAGING_PLATFORM_BASE_URL = "https://stagingplatform.comfy.org"
|
||||
|
||||
|
||||
def _is_staging_tier(host: str) -> bool:
|
||||
return host == _STAGING_API_HOST or host.endswith(_TESTENV_HOST_SUFFIX)
|
||||
|
||||
|
||||
def normalize_comfy_api_base(url: str) -> str:
|
||||
"""Rewrite a testenv's friendly main host to its comfy-api '-registry' sibling."""
|
||||
parsed = urlparse(url)
|
||||
host = parsed.hostname or ""
|
||||
if not host.endswith(_TESTENV_HOST_SUFFIX):
|
||||
return url
|
||||
label = host[: -len(_TESTENV_HOST_SUFFIX)]
|
||||
if label.endswith("-registry"):
|
||||
return url
|
||||
return f"{parsed.scheme or 'https'}://{label}-registry{_TESTENV_HOST_SUFFIX}"
|
||||
|
||||
|
||||
def environment_overrides_for_base(base_url: str) -> dict[str, Any] | None:
|
||||
"""The /features overrides for a staging-tier base, or None for prod."""
|
||||
if not _is_staging_tier(urlparse(base_url).hostname or ""):
|
||||
return None
|
||||
return {
|
||||
"comfy_api_base_url": normalize_comfy_api_base(base_url).rstrip("/"),
|
||||
"comfy_platform_base_url": _STAGING_PLATFORM_BASE_URL,
|
||||
"firebase_env": "dev",
|
||||
}
|
||||
|
||||
|
||||
def get_environment_overrides() -> dict[str, Any] | None:
|
||||
return environment_overrides_for_base(getattr(args, "comfy_api_base", "") or "")
|
||||
@@ -779,10 +779,6 @@ class ACEAudio(LatentFormat):
|
||||
latent_channels = 8
|
||||
latent_dimensions = 2
|
||||
|
||||
class SeedVR2(LatentFormat):
|
||||
latent_channels = 16
|
||||
latent_dimensions = 3
|
||||
|
||||
class ACEAudio15(LatentFormat):
|
||||
latent_channels = 64
|
||||
latent_dimensions = 1
|
||||
|
||||
@@ -15,24 +15,24 @@ def make_two_pass_attention(ar_len: int, transformer_options=None):
|
||||
The AR pass goes through SDPA directand bypasses wrappers, it is only ~1% of T at typical edit sizes.
|
||||
"""
|
||||
|
||||
def two_pass_attention(q, k, v, heads, enable_gqa=False, **kwargs):
|
||||
def two_pass_attention(q, k, v, heads, **kwargs):
|
||||
B, H, T, D = q.shape
|
||||
|
||||
if T < k.shape[2]: # KV-cache hot path: Q is shorter than K/V (cached AR prefix is in K/V only), all fresh Q positions are in the gen region, single full-attention call
|
||||
out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options, enable_gqa=enable_gqa)
|
||||
out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options)
|
||||
elif ar_len >= T:
|
||||
out = comfy.ops.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True, enable_gqa=enable_gqa)
|
||||
out = comfy.ops.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True)
|
||||
elif ar_len <= 0:
|
||||
out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options, enable_gqa=enable_gqa)
|
||||
out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options)
|
||||
else:
|
||||
out_ar = comfy.ops.scaled_dot_product_attention(
|
||||
q[:, :, :ar_len], k[:, :, :ar_len], v[:, :, :ar_len],
|
||||
attn_mask=None, dropout_p=0.0, is_causal=True, enable_gqa=enable_gqa,
|
||||
attn_mask=None, dropout_p=0.0, is_causal=True,
|
||||
)
|
||||
out_gen = optimized_attention(
|
||||
q[:, :, ar_len:], k, v, heads,
|
||||
mask=None, skip_reshape=True, skip_output_reshape=True,
|
||||
transformer_options=transformer_options, enable_gqa=enable_gqa,
|
||||
transformer_options=transformer_options,
|
||||
)
|
||||
out = torch.cat([out_ar, out_gen], dim=2)
|
||||
|
||||
|
||||
@@ -709,7 +709,7 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
|
||||
return out
|
||||
|
||||
try:
|
||||
@torch.library.custom_op("comfy::flash_attn", mutates_args=())
|
||||
@torch.library.custom_op("flash_attention::flash_attn", mutates_args=())
|
||||
def flash_attn_wrapper(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
|
||||
dropout_p: float = 0.0, causal: bool = False, softmax_scale: float = -1.0) -> torch.Tensor:
|
||||
softmax_scale_arg = None if softmax_scale == -1.0 else softmax_scale
|
||||
|
||||
@@ -22,7 +22,7 @@ def torch_cat_if_needed(xl, dim):
|
||||
else:
|
||||
return None
|
||||
|
||||
def get_timestep_embedding(timesteps, embedding_dim, flip_sin_to_cos=False, downscale_freq_shift=1):
|
||||
def get_timestep_embedding(timesteps, embedding_dim):
|
||||
"""
|
||||
This matches the implementation in Denoising Diffusion Probabilistic Models:
|
||||
From Fairseq.
|
||||
@@ -33,13 +33,11 @@ def get_timestep_embedding(timesteps, embedding_dim, flip_sin_to_cos=False, down
|
||||
assert len(timesteps.shape) == 1
|
||||
|
||||
half_dim = embedding_dim // 2
|
||||
emb = math.log(10000) / (half_dim - downscale_freq_shift)
|
||||
emb = math.log(10000) / (half_dim - 1)
|
||||
emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb)
|
||||
emb = emb.to(device=timesteps.device)
|
||||
emb = timesteps.float()[:, None] * emb[None, :]
|
||||
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
|
||||
if flip_sin_to_cos:
|
||||
emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
|
||||
if embedding_dim % 2 == 1: # zero pad
|
||||
emb = torch.nn.functional.pad(emb, (0,1,0,0))
|
||||
return emb
|
||||
|
||||
@@ -197,9 +197,6 @@ class PixDiT_T2I(nn.Module):
|
||||
"""Hook for subclasses to inject per-block state into the patch stream (e.g. PiD's LQ gate)."""
|
||||
return s
|
||||
|
||||
def _pre_pixel_blocks(self, s, **kwargs):
|
||||
return s
|
||||
|
||||
def _forward(self, x, timesteps, context=None, attention_mask=None, transformer_options={}, **kwargs):
|
||||
H_orig, W_orig = x.shape[2], x.shape[3]
|
||||
x = comfy.ldm.common_dit.pad_to_patch_size(x, (self.patch_size, self.patch_size))
|
||||
@@ -229,7 +226,6 @@ class PixDiT_T2I(nn.Module):
|
||||
s, y_emb = blk(s, y_emb, condition, pos_img, pos_txt, None, transformer_options=transformer_options)
|
||||
s = F.silu(t_emb + s)
|
||||
|
||||
s = self._pre_pixel_blocks(s, **kwargs)
|
||||
s_cond = s.view(B * L, self.hidden_size)
|
||||
x_pixels = self.pixel_embedder(x, patch_size=self.patch_size)
|
||||
for blk in self.pixel_blocks:
|
||||
|
||||
+14
-50
@@ -13,15 +13,15 @@ from .model import PixDiT_T2I
|
||||
from .modules import precompute_freqs_cis_2d
|
||||
|
||||
|
||||
class SigmaAwareGate(nn.Module):
|
||||
class SigmaAwareGatePerTokenPerDim(nn.Module):
|
||||
"""gate = sigmoid(content_proj(cat[x, lq]) - exp(log_alpha) * sigma); out = x + gate * lq.
|
||||
|
||||
Trained init gives ~0.88 gate at sigma=0, ~0.05 at sigma=1.
|
||||
"""
|
||||
|
||||
def __init__(self, dim: int, per_token: bool = False, dtype=None, device=None, operations=None):
|
||||
def __init__(self, dim: int, dtype=None, device=None, operations=None):
|
||||
super().__init__()
|
||||
self.content_proj = operations.Linear(dim * 2, 1 if per_token else dim, dtype=dtype, device=device)
|
||||
self.content_proj = operations.Linear(dim * 2, dim, dtype=dtype, device=device)
|
||||
self.log_alpha = nn.Parameter(torch.empty((), dtype=dtype, device=device))
|
||||
|
||||
def forward(self, x: torch.Tensor, lq: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor:
|
||||
@@ -36,15 +36,15 @@ class SigmaAwareGate(nn.Module):
|
||||
class ResBlock(nn.Module):
|
||||
"""Pre-activation ResNet block: GN -> SiLU -> Conv -> GN -> SiLU -> Conv + skip."""
|
||||
|
||||
def __init__(self, channels: int, num_groups: int = 4, conv_padding_mode: str = "zeros", dtype=None, device=None, operations=None):
|
||||
def __init__(self, channels: int, num_groups: int = 4, dtype=None, device=None, operations=None):
|
||||
super().__init__()
|
||||
self.block = nn.Sequential(
|
||||
operations.GroupNorm(num_groups, channels, dtype=dtype, device=device),
|
||||
nn.SiLU(),
|
||||
operations.Conv2d(channels, channels, kernel_size=3, padding=1, padding_mode=conv_padding_mode, dtype=dtype, device=device),
|
||||
operations.Conv2d(channels, channels, kernel_size=3, padding=1, dtype=dtype, device=device),
|
||||
operations.GroupNorm(num_groups, channels, dtype=dtype, device=device),
|
||||
nn.SiLU(),
|
||||
operations.Conv2d(channels, channels, kernel_size=3, padding=1, padding_mode=conv_padding_mode, dtype=dtype, device=device),
|
||||
operations.Conv2d(channels, channels, kernel_size=3, padding=1, dtype=dtype, device=device),
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
@@ -62,13 +62,9 @@ class LQProjection2D(nn.Module):
|
||||
patch_size: int = 16,
|
||||
sr_scale: int = 4,
|
||||
latent_spatial_down_factor: int = 8,
|
||||
latent_unpatchify_factor: int = 1,
|
||||
num_res_blocks: int = 4,
|
||||
num_outputs: int = 7,
|
||||
interval: int = 2,
|
||||
conv_padding_mode: str = "zeros",
|
||||
gate_per_token: bool = False,
|
||||
pit_output: bool = False,
|
||||
dtype=None, device=None, operations=None,
|
||||
):
|
||||
super().__init__()
|
||||
@@ -78,38 +74,34 @@ class LQProjection2D(nn.Module):
|
||||
self.patch_size = patch_size
|
||||
self.sr_scale = sr_scale
|
||||
self.latent_spatial_down_factor = latent_spatial_down_factor
|
||||
self.latent_unpatchify_factor = latent_unpatchify_factor
|
||||
self.num_outputs = num_outputs
|
||||
self.interval = interval
|
||||
|
||||
effective_latent_channels = latent_channels // (latent_unpatchify_factor * latent_unpatchify_factor)
|
||||
effective_spatial_down_factor = latent_spatial_down_factor // latent_unpatchify_factor
|
||||
z_to_patch_ratio = (sr_scale * effective_spatial_down_factor) / patch_size
|
||||
z_to_patch_ratio = (sr_scale * latent_spatial_down_factor) / patch_size
|
||||
self.z_to_patch_ratio = z_to_patch_ratio
|
||||
if z_to_patch_ratio >= 1:
|
||||
self.latent_fold_factor = 0
|
||||
latent_proj_in_ch = effective_latent_channels
|
||||
latent_proj_in_ch = latent_channels
|
||||
else:
|
||||
fold_factor = int(1 / z_to_patch_ratio)
|
||||
assert fold_factor * z_to_patch_ratio == 1.0
|
||||
self.latent_fold_factor = fold_factor
|
||||
latent_proj_in_ch = effective_latent_channels * fold_factor * fold_factor
|
||||
latent_proj_in_ch = latent_channels * fold_factor * fold_factor
|
||||
|
||||
layers = [
|
||||
operations.Conv2d(latent_proj_in_ch, hidden_dim, kernel_size=3, padding=1, padding_mode=conv_padding_mode, dtype=dtype, device=device),
|
||||
operations.Conv2d(latent_proj_in_ch, hidden_dim, kernel_size=3, padding=1, dtype=dtype, device=device),
|
||||
nn.SiLU(),
|
||||
operations.Conv2d(hidden_dim, hidden_dim, kernel_size=3, padding=1, padding_mode=conv_padding_mode, dtype=dtype, device=device),
|
||||
operations.Conv2d(hidden_dim, hidden_dim, kernel_size=3, padding=1, dtype=dtype, device=device),
|
||||
]
|
||||
for _ in range(num_res_blocks):
|
||||
layers.append(ResBlock(hidden_dim, conv_padding_mode=conv_padding_mode, dtype=dtype, device=device, operations=operations))
|
||||
layers.append(ResBlock(hidden_dim, dtype=dtype, device=device, operations=operations))
|
||||
self.latent_proj = nn.Sequential(*layers)
|
||||
|
||||
self.output_heads = nn.ModuleList(
|
||||
[operations.Linear(hidden_dim, out_dim, dtype=dtype, device=device) for _ in range(num_outputs)]
|
||||
)
|
||||
self.pit_head = operations.Linear(hidden_dim, out_dim, dtype=dtype, device=device) if pit_output else None
|
||||
self.gate_modules = nn.ModuleList(
|
||||
[SigmaAwareGate(out_dim, per_token=gate_per_token, dtype=dtype, device=device, operations=operations)
|
||||
[SigmaAwareGatePerTokenPerDim(out_dim, dtype=dtype, device=device, operations=operations)
|
||||
for _ in range(num_outputs)]
|
||||
)
|
||||
|
||||
@@ -123,11 +115,6 @@ class LQProjection2D(nn.Module):
|
||||
return self.gate_modules[out_idx](x, lq_feature, sigma)
|
||||
|
||||
def _align_latent_to_patch_grid(self, lq_latent: torch.Tensor, pH: int, pW: int) -> torch.Tensor:
|
||||
f = self.latent_unpatchify_factor
|
||||
if f > 1:
|
||||
B, C, H, W = lq_latent.shape
|
||||
lq_latent = lq_latent.reshape(B, C // (f * f), f, f, H, W)
|
||||
lq_latent = lq_latent.permute(0, 1, 4, 2, 5, 3).reshape(B, C // (f * f), H * f, W * f)
|
||||
B, z_dim = lq_latent.shape[:2]
|
||||
if self.z_to_patch_ratio >= 1:
|
||||
if lq_latent.shape[2] != pH or lq_latent.shape[3] != pW:
|
||||
@@ -147,10 +134,7 @@ class LQProjection2D(nn.Module):
|
||||
feat = self._align_latent_to_patch_grid(lq_latent, target_pH, target_pW)
|
||||
B, C, H, W = feat.shape
|
||||
tokens = feat.permute(0, 2, 3, 1).contiguous().view(B, H * W, C)
|
||||
outputs = [head(tokens) for head in self.output_heads]
|
||||
if self.pit_head is not None:
|
||||
outputs.append(self.pit_head(tokens))
|
||||
return outputs
|
||||
return [head(tokens) for head in self.output_heads]
|
||||
|
||||
|
||||
class PidNet(PixDiT_T2I):
|
||||
@@ -164,10 +148,6 @@ class PidNet(PixDiT_T2I):
|
||||
lq_interval: int = 2,
|
||||
sr_scale: int = 4,
|
||||
latent_spatial_down_factor: int = 8,
|
||||
lq_latent_unpatchify_factor: int = 1,
|
||||
lq_conv_padding_mode: str = "zeros",
|
||||
lq_gate_per_token: bool = False,
|
||||
pit_lq_inject: bool = False,
|
||||
rope_ref_h: int = 1024, # NTK ref resolution in PIXEL units: 1024px / patch=16 -> grid_ref=64.
|
||||
rope_ref_w: int = 1024,
|
||||
image_model=None,
|
||||
@@ -185,8 +165,6 @@ class PidNet(PixDiT_T2I):
|
||||
for blk in self.pixel_blocks:
|
||||
blk._rope_fn = _pit_rope_fn
|
||||
|
||||
self.pit_lq_inject = pit_lq_inject
|
||||
|
||||
num_lq_outputs = (self.patch_depth + lq_interval - 1) // lq_interval
|
||||
self.lq_proj = LQProjection2D(
|
||||
latent_channels=lq_latent_channels,
|
||||
@@ -195,20 +173,13 @@ class PidNet(PixDiT_T2I):
|
||||
patch_size=self.patch_size,
|
||||
sr_scale=sr_scale,
|
||||
latent_spatial_down_factor=latent_spatial_down_factor,
|
||||
latent_unpatchify_factor=lq_latent_unpatchify_factor,
|
||||
num_res_blocks=lq_num_res_blocks,
|
||||
num_outputs=num_lq_outputs,
|
||||
interval=lq_interval,
|
||||
conv_padding_mode=lq_conv_padding_mode,
|
||||
gate_per_token=lq_gate_per_token,
|
||||
pit_output=pit_lq_inject,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
operations=operations,
|
||||
)
|
||||
self.pit_lq_gate = SigmaAwareGate(
|
||||
self.hidden_size, per_token=lq_gate_per_token, dtype=dtype, device=device, operations=operations
|
||||
) if pit_lq_inject else None
|
||||
|
||||
def _fetch_patch_pos(self, height, width, device, dtype, **rope_opts):
|
||||
return precompute_freqs_cis_2d(
|
||||
@@ -226,11 +197,6 @@ class PidNet(PixDiT_T2I):
|
||||
return s
|
||||
return self.lq_proj.gate(s, pid_lq_features[out_idx], pid_degrade_sigma, out_idx)
|
||||
|
||||
def _pre_pixel_blocks(self, s, pid_pit_lq_feature=None, pid_degrade_sigma=None, **kwargs):
|
||||
if pid_pit_lq_feature is None:
|
||||
return s
|
||||
return self.pit_lq_gate(s, pid_pit_lq_feature, pid_degrade_sigma)
|
||||
|
||||
def _forward(self, x, timesteps, context=None, attention_mask=None, transformer_options={}, lq_latent=None, degrade_sigma=None, **kwargs):
|
||||
if lq_latent is None:
|
||||
raise ValueError("PidNet requires lq_latent — attach via PiDConditioning")
|
||||
@@ -250,14 +216,12 @@ class PidNet(PixDiT_T2I):
|
||||
degrade_sigma = degrade_sigma.expand(B).contiguous()
|
||||
|
||||
lq_features = self.lq_proj(lq_latent=lq_latent.to(x), target_pH=Hs, target_pW=Ws)
|
||||
pit_lq_feature = lq_features.pop() if self.pit_lq_inject else None
|
||||
|
||||
return super()._forward(
|
||||
x, timesteps,
|
||||
context=context, attention_mask=attention_mask,
|
||||
transformer_options=transformer_options,
|
||||
pid_lq_features=lq_features,
|
||||
pid_pit_lq_feature=pit_lq_feature,
|
||||
pid_degrade_sigma=degrade_sigma,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
@@ -1,51 +0,0 @@
|
||||
import torch
|
||||
|
||||
from comfy.ldm.modules import attention as _attention
|
||||
|
||||
|
||||
def _var_attention_qkv(q, k, v, heads, skip_reshape):
|
||||
if skip_reshape:
|
||||
return q, k, v, q.shape[-1]
|
||||
total_tokens, embed_dim = q.shape
|
||||
head_dim = embed_dim // heads
|
||||
return (
|
||||
q.view(total_tokens, heads, head_dim),
|
||||
k.view(k.shape[0], heads, head_dim),
|
||||
v.view(v.shape[0], heads, head_dim),
|
||||
head_dim,
|
||||
)
|
||||
|
||||
|
||||
def _var_attention_output(out, heads, head_dim, skip_output_reshape):
|
||||
if skip_output_reshape:
|
||||
return out
|
||||
return out.reshape(-1, heads * head_dim)
|
||||
|
||||
|
||||
def var_attention_optimized_split(q, k, v, heads, cu_seqlens_q, cu_seqlens_k, *args, skip_reshape=False, skip_output_reshape=False, **kwargs):
|
||||
q, k, v, head_dim = _var_attention_qkv(q, k, v, heads, skip_reshape)
|
||||
|
||||
q_split_indices = cu_seqlens_q[1:-1]
|
||||
k_split_indices = cu_seqlens_k[1:-1]
|
||||
if k.shape[0] != v.shape[0]:
|
||||
raise ValueError("cu_seqlens_k does not match v token count")
|
||||
|
||||
q_splits = torch.tensor_split(q, q_split_indices, dim=0)
|
||||
k_splits = torch.tensor_split(k, k_split_indices, dim=0)
|
||||
v_splits = torch.tensor_split(v, k_split_indices, dim=0)
|
||||
if len(q_splits) != len(k_splits) or len(q_splits) != len(v_splits):
|
||||
raise ValueError("cu_seqlens_q and cu_seqlens_k must describe the same sequence count")
|
||||
|
||||
out = []
|
||||
for q_i, k_i, v_i in zip(q_splits, k_splits, v_splits):
|
||||
q_i = q_i.permute(1, 0, 2).unsqueeze(0)
|
||||
k_i = k_i.permute(1, 0, 2).unsqueeze(0)
|
||||
v_i = v_i.permute(1, 0, 2).unsqueeze(0)
|
||||
out_i = _attention.optimized_attention(q_i, k_i, v_i, heads, skip_reshape=True, skip_output_reshape=True)
|
||||
out.append(out_i.squeeze(0).permute(1, 0, 2))
|
||||
|
||||
out = torch.cat(out, dim=0)
|
||||
return _var_attention_output(out, heads, head_dim, skip_output_reshape)
|
||||
|
||||
|
||||
optimized_var_attention = var_attention_optimized_split
|
||||
@@ -1,301 +0,0 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
|
||||
from comfy.ldm.seedvr.constants import (
|
||||
CIELAB_DELTA,
|
||||
CIELAB_KAPPA,
|
||||
D65_WHITE_X,
|
||||
D65_WHITE_Z,
|
||||
WAVELET_DECOMP_LEVELS,
|
||||
)
|
||||
|
||||
|
||||
def wavelet_blur(image: Tensor, radius):
|
||||
max_safe_radius = max(1, min(image.shape[-2:]) // 8)
|
||||
if radius > max_safe_radius:
|
||||
radius = max_safe_radius
|
||||
|
||||
num_channels = image.shape[1]
|
||||
|
||||
kernel_vals = [
|
||||
[0.0625, 0.125, 0.0625],
|
||||
[0.125, 0.25, 0.125],
|
||||
[0.0625, 0.125, 0.0625],
|
||||
]
|
||||
kernel = torch.tensor(kernel_vals, dtype=image.dtype, device=image.device)
|
||||
kernel = kernel[None, None].repeat(num_channels, 1, 1, 1)
|
||||
|
||||
image = F.pad(image, (radius, radius, radius, radius), mode='replicate')
|
||||
output = F.conv2d(image, kernel, groups=num_channels, dilation=radius)
|
||||
|
||||
return output
|
||||
|
||||
def wavelet_decomposition(image: Tensor, levels: int = WAVELET_DECOMP_LEVELS):
|
||||
high_freq = torch.zeros_like(image)
|
||||
|
||||
for i in range(levels):
|
||||
radius = 2 ** i
|
||||
low_freq = wavelet_blur(image, radius)
|
||||
high_freq.add_(image).sub_(low_freq)
|
||||
image = low_freq
|
||||
|
||||
return high_freq, low_freq
|
||||
|
||||
def wavelet_reconstruction(content_feat: Tensor, style_feat: Tensor) -> Tensor:
|
||||
|
||||
if content_feat.shape != style_feat.shape:
|
||||
if len(content_feat.shape) >= 3:
|
||||
style_feat = F.interpolate(
|
||||
style_feat,
|
||||
size=content_feat.shape[-2:],
|
||||
mode='bilinear',
|
||||
align_corners=False
|
||||
)
|
||||
|
||||
content_high_freq, content_low_freq = wavelet_decomposition(content_feat)
|
||||
del content_low_freq
|
||||
|
||||
style_high_freq, style_low_freq = wavelet_decomposition(style_feat)
|
||||
del style_high_freq
|
||||
|
||||
if content_high_freq.shape != style_low_freq.shape:
|
||||
style_low_freq = F.interpolate(
|
||||
style_low_freq,
|
||||
size=content_high_freq.shape[-2:],
|
||||
mode='bilinear',
|
||||
align_corners=False
|
||||
)
|
||||
|
||||
content_high_freq.add_(style_low_freq)
|
||||
|
||||
return content_high_freq.clamp_(-1.0, 1.0)
|
||||
|
||||
def _histogram_matching_channel(source: Tensor, reference: Tensor) -> Tensor:
|
||||
original_shape = source.shape
|
||||
|
||||
source_flat = source.flatten()
|
||||
reference_flat = reference.flatten()
|
||||
|
||||
source_sorted, source_indices = torch.sort(source_flat)
|
||||
reference_sorted, _ = torch.sort(reference_flat)
|
||||
del reference_flat
|
||||
|
||||
n_source = len(source_sorted)
|
||||
n_reference = len(reference_sorted)
|
||||
|
||||
if n_source == n_reference:
|
||||
matched_sorted = reference_sorted
|
||||
else:
|
||||
source_quantiles = torch.linspace(0, 1, n_source, device=source.device)
|
||||
ref_indices = (source_quantiles * (n_reference - 1)).long()
|
||||
ref_indices.clamp_(0, n_reference - 1)
|
||||
matched_sorted = reference_sorted[ref_indices]
|
||||
del source_quantiles, ref_indices, reference_sorted
|
||||
|
||||
del source_sorted, source_flat
|
||||
|
||||
inverse_indices = torch.argsort(source_indices)
|
||||
del source_indices
|
||||
matched_flat = matched_sorted[inverse_indices]
|
||||
del matched_sorted, inverse_indices
|
||||
|
||||
return matched_flat.reshape(original_shape)
|
||||
|
||||
def _lab_to_rgb_batch(lab: Tensor, matrix_inv: Tensor, epsilon: float, kappa: float) -> Tensor:
|
||||
L, a, b = lab[:, 0], lab[:, 1], lab[:, 2]
|
||||
|
||||
fy = (L + 16.0) / 116.0
|
||||
fx = a.div(500.0).add_(fy)
|
||||
fz = fy - b / 200.0
|
||||
del L, a, b
|
||||
|
||||
x = torch.where(
|
||||
fx > epsilon,
|
||||
torch.pow(fx, 3.0),
|
||||
fx.mul(116.0).sub_(16.0).div_(kappa)
|
||||
)
|
||||
y = torch.where(
|
||||
fy > epsilon,
|
||||
torch.pow(fy, 3.0),
|
||||
fy.mul(116.0).sub_(16.0).div_(kappa)
|
||||
)
|
||||
z = torch.where(
|
||||
fz > epsilon,
|
||||
torch.pow(fz, 3.0),
|
||||
fz.mul(116.0).sub_(16.0).div_(kappa)
|
||||
)
|
||||
del fx, fy, fz
|
||||
|
||||
x.mul_(D65_WHITE_X)
|
||||
z.mul_(D65_WHITE_Z)
|
||||
|
||||
xyz = torch.stack([x, y, z], dim=1)
|
||||
del x, y, z
|
||||
|
||||
B, _, H, W = xyz.shape
|
||||
xyz_flat = xyz.permute(0, 2, 3, 1).reshape(-1, 3)
|
||||
del xyz
|
||||
|
||||
xyz_flat = xyz_flat.to(dtype=matrix_inv.dtype)
|
||||
rgb_linear_flat = torch.matmul(xyz_flat, matrix_inv.T)
|
||||
del xyz_flat
|
||||
|
||||
rgb_linear = rgb_linear_flat.reshape(B, H, W, 3).permute(0, 3, 1, 2)
|
||||
del rgb_linear_flat
|
||||
|
||||
mask = rgb_linear > 0.0031308
|
||||
rgb = torch.where(
|
||||
mask,
|
||||
torch.pow(torch.clamp(rgb_linear, min=0.0), 1.0 / 2.4).mul_(1.055).sub_(0.055),
|
||||
rgb_linear * 12.92
|
||||
)
|
||||
del mask, rgb_linear
|
||||
|
||||
return torch.clamp(rgb, 0.0, 1.0)
|
||||
|
||||
def _rgb_to_lab_batch(rgb: Tensor, matrix: Tensor, epsilon: float, kappa: float) -> Tensor:
|
||||
mask = rgb > 0.04045
|
||||
rgb_linear = torch.where(
|
||||
mask,
|
||||
torch.pow((rgb + 0.055) / 1.055, 2.4),
|
||||
rgb / 12.92
|
||||
)
|
||||
del mask
|
||||
|
||||
B, _, H, W = rgb_linear.shape
|
||||
rgb_flat = rgb_linear.permute(0, 2, 3, 1).reshape(-1, 3)
|
||||
del rgb_linear
|
||||
|
||||
rgb_flat = rgb_flat.to(dtype=matrix.dtype)
|
||||
xyz_flat = torch.matmul(rgb_flat, matrix.T)
|
||||
del rgb_flat
|
||||
|
||||
xyz = xyz_flat.reshape(B, H, W, 3).permute(0, 3, 1, 2)
|
||||
del xyz_flat
|
||||
|
||||
xyz[:, 0].div_(D65_WHITE_X)
|
||||
xyz[:, 2].div_(D65_WHITE_Z)
|
||||
|
||||
epsilon_cubed = epsilon ** 3
|
||||
mask = xyz > epsilon_cubed
|
||||
f_xyz = torch.where(
|
||||
mask,
|
||||
torch.pow(xyz, 1.0 / 3.0),
|
||||
xyz.mul(kappa).add_(16.0).div_(116.0)
|
||||
)
|
||||
del xyz, mask
|
||||
|
||||
L = f_xyz[:, 1].mul(116.0).sub_(16.0)
|
||||
a = (f_xyz[:, 0] - f_xyz[:, 1]).mul_(500.0)
|
||||
b = (f_xyz[:, 1] - f_xyz[:, 2]).mul_(200.0)
|
||||
del f_xyz
|
||||
|
||||
return torch.stack([L, a, b], dim=1)
|
||||
|
||||
def lab_color_transfer(
|
||||
content_feat: Tensor,
|
||||
style_feat: Tensor,
|
||||
luminance_weight: float = 0.8
|
||||
) -> Tensor:
|
||||
content_feat = wavelet_reconstruction(content_feat, style_feat)
|
||||
|
||||
if content_feat.shape != style_feat.shape:
|
||||
style_feat = F.interpolate(
|
||||
style_feat,
|
||||
size=content_feat.shape[-2:],
|
||||
mode='bilinear',
|
||||
align_corners=False
|
||||
)
|
||||
|
||||
device = content_feat.device
|
||||
original_dtype = content_feat.dtype
|
||||
content_feat = content_feat.float()
|
||||
style_feat = style_feat.float()
|
||||
|
||||
rgb_to_xyz_matrix = torch.tensor([
|
||||
[0.4124564, 0.3575761, 0.1804375],
|
||||
[0.2126729, 0.7151522, 0.0721750],
|
||||
[0.0193339, 0.1191920, 0.9503041]
|
||||
], dtype=torch.float32, device=device)
|
||||
|
||||
xyz_to_rgb_matrix = torch.tensor([
|
||||
[ 3.2404542, -1.5371385, -0.4985314],
|
||||
[-0.9692660, 1.8760108, 0.0415560],
|
||||
[ 0.0556434, -0.2040259, 1.0572252]
|
||||
], dtype=torch.float32, device=device)
|
||||
|
||||
epsilon = CIELAB_DELTA
|
||||
kappa = CIELAB_KAPPA
|
||||
|
||||
content_feat.add_(1.0).mul_(0.5).clamp_(0.0, 1.0)
|
||||
style_feat.add_(1.0).mul_(0.5).clamp_(0.0, 1.0)
|
||||
|
||||
content_lab = _rgb_to_lab_batch(content_feat, rgb_to_xyz_matrix, epsilon, kappa)
|
||||
del content_feat
|
||||
|
||||
style_lab = _rgb_to_lab_batch(style_feat, rgb_to_xyz_matrix, epsilon, kappa)
|
||||
del style_feat, rgb_to_xyz_matrix
|
||||
|
||||
matched_a = _histogram_matching_channel(content_lab[:, 1], style_lab[:, 1])
|
||||
matched_b = _histogram_matching_channel(content_lab[:, 2], style_lab[:, 2])
|
||||
|
||||
if luminance_weight < 1.0:
|
||||
matched_L = _histogram_matching_channel(content_lab[:, 0], style_lab[:, 0])
|
||||
result_L = content_lab[:, 0].mul(luminance_weight).add_(matched_L.mul(1.0 - luminance_weight))
|
||||
del matched_L
|
||||
else:
|
||||
result_L = content_lab[:, 0]
|
||||
|
||||
del content_lab, style_lab
|
||||
|
||||
result_lab = torch.stack([result_L, matched_a, matched_b], dim=1)
|
||||
del result_L, matched_a, matched_b
|
||||
|
||||
result_rgb = _lab_to_rgb_batch(result_lab, xyz_to_rgb_matrix, epsilon, kappa)
|
||||
del result_lab, xyz_to_rgb_matrix
|
||||
|
||||
result = result_rgb.mul_(2.0).sub_(1.0)
|
||||
del result_rgb
|
||||
|
||||
result = result.to(original_dtype)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def wavelet_color_transfer(content_feat: Tensor, style_feat: Tensor) -> Tensor:
|
||||
return wavelet_reconstruction(content_feat, style_feat)
|
||||
|
||||
|
||||
def adain_color_transfer(content_feat: Tensor, style_feat: Tensor, eps: float = 1e-5) -> Tensor:
|
||||
if content_feat.shape != style_feat.shape:
|
||||
style_feat = F.interpolate(
|
||||
style_feat,
|
||||
size=content_feat.shape[-2:],
|
||||
mode='bilinear',
|
||||
align_corners=False,
|
||||
)
|
||||
|
||||
original_dtype = content_feat.dtype
|
||||
content_feat = content_feat.float()
|
||||
style_feat = style_feat.float()
|
||||
|
||||
b, c = content_feat.shape[:2]
|
||||
content_flat = content_feat.reshape(b, c, -1)
|
||||
style_flat = style_feat.reshape(b, c, -1)
|
||||
|
||||
content_mean = content_flat.mean(dim=2).reshape(b, c, 1, 1)
|
||||
content_std = (content_flat.var(dim=2, correction=0) + eps).sqrt().reshape(b, c, 1, 1)
|
||||
style_mean = style_flat.mean(dim=2).reshape(b, c, 1, 1)
|
||||
style_std = (style_flat.var(dim=2, correction=0) + eps).sqrt().reshape(b, c, 1, 1)
|
||||
del content_flat, style_flat
|
||||
|
||||
normalized = (content_feat - content_mean) / content_std
|
||||
del content_mean, content_std
|
||||
result = normalized * style_std + style_mean
|
||||
del normalized, style_mean, style_std
|
||||
|
||||
result = result.clamp_(-1.0, 1.0)
|
||||
if result.dtype != original_dtype:
|
||||
result = result.to(original_dtype)
|
||||
return result
|
||||
@@ -1,48 +0,0 @@
|
||||
"""SeedVR2 constants."""
|
||||
|
||||
# Temporal chunk-size law: the sampler's activation wall is linear in
|
||||
# T_latent * pixel area (17-cell resolution sweep + T bisection, RTX 5090, 3b fp16):
|
||||
# max_latent_frames = (free_GiB - RESERVED - K*SIGMA) / (GIB_PER_MPX_FRAME * megapixels)
|
||||
# RESERVED covers model staging plus fixed CUDA/torch overhead; SIGMA is the measured
|
||||
# run-to-run spread of the wall; K=4 trades ~10% smaller chunks for ~1e-5 OOM odds.
|
||||
SEEDVR2_CHUNK_GIB_PER_MPX_FRAME = 0.55
|
||||
SEEDVR2_CHUNK_RESERVED_GIB = 8.5
|
||||
SEEDVR2_CHUNK_SIGMA_GIB = 0.55
|
||||
SEEDVR2_CHUNK_SIGMA_K = 4
|
||||
|
||||
SEEDVR2_7B_VID_DIM = 3072
|
||||
SEEDVR2_OOM_BACKOFF_DIVISOR = 2
|
||||
SEEDVR2_DTYPE_BYTES_FLOOR = 4
|
||||
SEEDVR2_7B_MLP_CHUNK = 8192
|
||||
SEEDVR2_ROPE_PARTIAL_CHUNK_TOKENS = 4096 # partial-RoPE application token-chunk.
|
||||
SEEDVR2_LATENT_CHANNELS = 16
|
||||
|
||||
SEEDVR2_COLOR_MEM_HEADROOM = 0.75
|
||||
SEEDVR2_LAB_SCALE_MULTIPLIER = 13
|
||||
SEEDVR2_WAVELET_SCALE_MULTIPLIER = 10 # per-frame byte multiplier, wavelet path.
|
||||
SEEDVR2_ADAIN_SCALE_MULTIPLIER = 6
|
||||
|
||||
BYTEDANCE_VAE_SCALING_FACTOR = 0.9152 # configs_3b/main.yaml:57.
|
||||
BYTEDANCE_VAE_SHIFTING_FACTOR = 0.0
|
||||
BYTEDANCE_VAE_CONV_MEM_GIB = 0.5
|
||||
BYTEDANCE_VAE_NORM_MEM_GIB = 0.5
|
||||
BYTEDANCE_LOGVAR_CLAMP_MIN = -30.0 # video_vae_v3/modules/types.py:28.
|
||||
BYTEDANCE_LOGVAR_CLAMP_MAX = 20.0 # video_vae_v3/modules/types.py:28.
|
||||
BYTEDANCE_GN_CHUNKS_FP16 = 4 # causal_inflation_lib.py:351 (GroupNorm chunk count, fp16).
|
||||
BYTEDANCE_GN_CHUNKS_FP32 = 2 # causal_inflation_lib.py:351 (GroupNorm chunk count, fp32).
|
||||
BYTEDANCE_BLOCK_OUT_CHANNELS = (128, 256, 512, 512) # s8_c16_t4_inflation_sd3.yaml:7-11.
|
||||
BYTEDANCE_SLICING_SAMPLE_MIN = 4 # s8_c16_t4_inflation_sd3.yaml:22 (slicing_sample_min_size).
|
||||
BYTEDANCE_VAE_TEMPORAL_DOWNSAMPLE = 4 # infer.py:230 (temporal_downsample_factor); the 4n+1 factor.
|
||||
BYTEDANCE_VAE_SPATIAL_DOWNSAMPLE = 8 # infer.py:231 (spatial_downsample_factor).
|
||||
BYTEDANCE_720P_REF_AREA = 45 * 80 # dit_v2/window.py:32 (720p reference area for window scaling).
|
||||
BYTEDANCE_MAX_TEMPORAL_WINDOW = 30 # dit_v2/window.py:35 (max temporal window frames).
|
||||
BYTEDANCE_ROPE_MAX_FREQ = 256 # dit_v2/rope.py:31 (pixel-RoPE max frequency).
|
||||
BYTEDANCE_SINUSOIDAL_DIM = 256 # dit_3b/nadit.py:120 (timestep sinusoidal embed dim).
|
||||
|
||||
ROPE_THETA = 10000 # RoPE base; Su et al., "RoFormer", arXiv:2104.09864.
|
||||
|
||||
CIELAB_DELTA = 6.0 / 29.0 # CIE 15 (delta).
|
||||
CIELAB_KAPPA = (29.0 / 3.0) ** 3 # CIE 15 (kappa).
|
||||
D65_WHITE_X = 0.95047 # CIE D65 standard illuminant Xn (Yn = 1).
|
||||
D65_WHITE_Z = 1.08883 # CIE D65 standard illuminant Zn.
|
||||
WAVELET_DECOMP_LEVELS = 5 # wavelet color-fix decomposition depth (GIMP/Krita; StableSR).
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -55,7 +55,6 @@ import comfy.ldm.pixeldit.model
|
||||
import comfy.ldm.pixeldit.pid
|
||||
import comfy.ldm.ace.model
|
||||
import comfy.ldm.omnigen.omnigen2
|
||||
import comfy.ldm.seedvr.model
|
||||
import comfy.ldm.boogu.model
|
||||
import comfy.ldm.qwen_image.model
|
||||
import comfy.ldm.ideogram4.model
|
||||
@@ -933,17 +932,6 @@ class HunyuanDiT(BaseModel):
|
||||
out['image_meta_size'] = comfy.conds.CONDRegular(torch.FloatTensor([[height, width, target_height, target_width, 0, 0]]))
|
||||
return out
|
||||
|
||||
class SeedVR2(BaseModel):
|
||||
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
|
||||
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.seedvr.model.NaDiT)
|
||||
|
||||
def extra_conds(self, **kwargs):
|
||||
out = super().extra_conds(**kwargs)
|
||||
condition = kwargs.get("condition", None)
|
||||
if condition is not None:
|
||||
out["condition"] = comfy.conds.CONDRegular(condition)
|
||||
return out
|
||||
|
||||
class PixArt(BaseModel):
|
||||
def __init__(self, model_config, model_type=ModelType.EPS, device=None):
|
||||
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.pixart.pixartms.PixArtMS)
|
||||
|
||||
@@ -470,46 +470,15 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
|
||||
# PiD (Pixel Diffusion Decoder). Must check BEFORE plain PixelDiT_T2I.
|
||||
_lq_w_key = '{}lq_proj.latent_proj.0.weight'.format(key_prefix)
|
||||
if _lq_w_key in state_dict_keys:
|
||||
latent_proj_in_channels = int(state_dict[_lq_w_key].shape[1])
|
||||
hidden_dim = int(state_dict[_lq_w_key].shape[0])
|
||||
in_ch = int(state_dict[_lq_w_key].shape[1])
|
||||
_gate_prefix = '{}lq_proj.gate_modules.'.format(key_prefix)
|
||||
num_gates = len({k[len(_gate_prefix):].split('.')[0]
|
||||
for k in state_dict_keys if k.startswith(_gate_prefix)})
|
||||
pid_v1_5 = '{}lq_proj.pit_head.weight'.format(key_prefix) in state_dict_keys
|
||||
dit_config = {"image_model": "pid",
|
||||
"lq_hidden_dim": hidden_dim}
|
||||
"lq_latent_channels": in_ch,
|
||||
"latent_spatial_down_factor": 16 if in_ch >= 64 else 8}
|
||||
if num_gates > 0:
|
||||
dit_config["lq_interval"] = (14 + num_gates - 1) // num_gates
|
||||
if pid_v1_5:
|
||||
pid_v1_5_variants = {
|
||||
16: { # Flux and QwenImage
|
||||
"lq_latent_channels": 16,
|
||||
"latent_spatial_down_factor": 8,
|
||||
"lq_latent_unpatchify_factor": 1,
|
||||
},
|
||||
32: { # Flux2 after 2x latent unpatchify
|
||||
"lq_latent_channels": 128,
|
||||
"latent_spatial_down_factor": 16,
|
||||
"lq_latent_unpatchify_factor": 2,
|
||||
},
|
||||
}
|
||||
variant = pid_v1_5_variants.get(latent_proj_in_channels)
|
||||
if variant is None:
|
||||
raise ValueError(f"Unsupported PiD v1.5 latent projection with {latent_proj_in_channels} input channels")
|
||||
gate_weight = state_dict['{}lq_proj.gate_modules.0.content_proj.weight'.format(key_prefix)]
|
||||
dit_config.update(variant)
|
||||
dit_config.update({
|
||||
"lq_conv_padding_mode": "replicate",
|
||||
"lq_gate_per_token": gate_weight.shape[0] == 1,
|
||||
"pit_lq_inject": True,
|
||||
"rope_ref_h": 2048,
|
||||
"rope_ref_w": 2048,
|
||||
})
|
||||
else:
|
||||
dit_config.update({
|
||||
"lq_latent_channels": latent_proj_in_channels,
|
||||
"latent_spatial_down_factor": 16 if latent_proj_in_channels >= 64 else 8,
|
||||
})
|
||||
return dit_config
|
||||
|
||||
if '{}core.pixel_embedder.proj.weight'.format(key_prefix) in state_dict_keys: # PixelDiT T2I
|
||||
@@ -629,44 +598,6 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
|
||||
|
||||
return dit_config
|
||||
|
||||
seedvr2_7b_separate_key = "{}blocks.35.mlp.vid.proj_out.weight".format(key_prefix)
|
||||
if seedvr2_7b_separate_key in state_dict_keys and state_dict[seedvr2_7b_separate_key].shape[0] == 3072: # seedvr2 7b
|
||||
dit_config = {}
|
||||
dit_config["image_model"] = "seedvr2"
|
||||
dit_config["vid_dim"] = 3072
|
||||
dit_config["heads"] = 24
|
||||
dit_config["num_layers"] = 36
|
||||
# This checkpoint uses separate vid/txt MMModule keys in every block.
|
||||
dit_config["mm_layers"] = 36
|
||||
dit_config["norm_eps"] = 1e-5
|
||||
dit_config["rope_type"] = "rope3d"
|
||||
dit_config["rope_dim"] = 64
|
||||
dit_config["mlp_type"] = "normal"
|
||||
return dit_config
|
||||
if "{}blocks.35.mlp.all.proj_in_gate.weight".format(key_prefix) in state_dict_keys: # seedvr2 7b
|
||||
dit_config = {}
|
||||
dit_config["image_model"] = "seedvr2"
|
||||
dit_config["vid_dim"] = 3072
|
||||
dit_config["heads"] = 24
|
||||
dit_config["num_layers"] = 36
|
||||
# This checkpoint uses shared all.* MMModule keys after the initial blocks.
|
||||
dit_config["mm_layers"] = 10
|
||||
dit_config["norm_eps"] = 1e-5
|
||||
dit_config["rope_type"] = "rope3d"
|
||||
dit_config["rope_dim"] = 64
|
||||
dit_config["mlp_type"] = "swiglu"
|
||||
return dit_config
|
||||
if "{}blocks.31.mlp.all.proj_in_gate.weight".format(key_prefix) in state_dict_keys: # seedvr2 3b
|
||||
dit_config = {}
|
||||
dit_config["image_model"] = "seedvr2"
|
||||
dit_config["vid_dim"] = 2560
|
||||
dit_config["heads"] = 20
|
||||
dit_config["num_layers"] = 32
|
||||
dit_config["norm_eps"] = 1.0e-05
|
||||
dit_config["mlp_type"] = "swiglu"
|
||||
dit_config["vid_out_norm"] = True
|
||||
return dit_config
|
||||
|
||||
if '{}head.modulation'.format(key_prefix) in state_dict_keys: # Wan 2.1
|
||||
dit_config = {}
|
||||
dit_config["image_model"] = "wan2.1"
|
||||
@@ -1188,10 +1119,9 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
|
||||
|
||||
return unet_config
|
||||
|
||||
|
||||
def model_config_from_unet_config(unet_config, state_dict=None, unet_key_prefix=""):
|
||||
def model_config_from_unet_config(unet_config, state_dict=None):
|
||||
for model_config in comfy.supported_models.models:
|
||||
if model_config.matches(unet_config, state_dict, unet_key_prefix=unet_key_prefix):
|
||||
if model_config.matches(unet_config, state_dict):
|
||||
return model_config(unet_config)
|
||||
|
||||
logging.error("no match {}".format(unet_config))
|
||||
@@ -1201,7 +1131,7 @@ def model_config_from_unet(state_dict, unet_key_prefix, use_base_if_no_match=Fal
|
||||
unet_config = detect_unet_config(state_dict, unet_key_prefix, metadata=metadata)
|
||||
if unet_config is None:
|
||||
return None
|
||||
model_config = model_config_from_unet_config(unet_config, state_dict, unet_key_prefix)
|
||||
model_config = model_config_from_unet_config(unet_config, state_dict)
|
||||
if model_config is None and use_base_if_no_match:
|
||||
model_config = comfy.supported_models_base.BASE(unet_config)
|
||||
|
||||
|
||||
@@ -616,8 +616,6 @@ PIN_PRESSURE_HYSTERESIS = 256 * 1024 * 1024
|
||||
#Freeing registerables on pressure does imply a GPU sync, so go big on
|
||||
#the hysteresis so each expensive sync gives us back a good chunk.
|
||||
REGISTERABLE_PIN_HYSTERESIS = 2048 * 1024 * 1024
|
||||
WINDOWS_PIN_EVICTION_SWAP_PERCENT = 5.0
|
||||
WINDOWS_PIN_EVICTION_EMERGENCY_AVAILABLE = 512 * 1024 ** 2
|
||||
|
||||
def module_size(module):
|
||||
module_mem = 0
|
||||
@@ -644,15 +642,6 @@ def free_pins(size, evict_active=False):
|
||||
size -= freed
|
||||
return freed_total
|
||||
|
||||
def should_free_pins_for_ram_pressure(shortfall):
|
||||
if shortfall <= 0:
|
||||
return False
|
||||
if not WINDOWS:
|
||||
return True
|
||||
if psutil.virtual_memory().available < WINDOWS_PIN_EVICTION_EMERGENCY_AVAILABLE:
|
||||
return True
|
||||
return psutil.swap_memory().percent >= WINDOWS_PIN_EVICTION_SWAP_PERCENT
|
||||
|
||||
def ensure_pin_budget(size, evict_active=False):
|
||||
if args.high_ram:
|
||||
return True
|
||||
|
||||
+7
-33
@@ -1104,21 +1104,6 @@ def _load_quantized_module(module, super_load, state_dict, prefix, local_metadat
|
||||
scales["convrot_groupsize"] = int(
|
||||
layer_conf.get("convrot_groupsize", params_conf.get("convrot_groupsize", 256))
|
||||
)
|
||||
elif module.quant_format == "convrot_w4a4":
|
||||
scale = pop_scale("weight_scale")
|
||||
if scale is None:
|
||||
raise ValueError(f"Missing ConvRot W4A4 weight scale for layer {layer_name}")
|
||||
params_conf = layer_conf.get("params", {})
|
||||
if not isinstance(params_conf, dict):
|
||||
params_conf = {}
|
||||
scales = {
|
||||
"scale": scale,
|
||||
"convrot_groupsize": int(
|
||||
layer_conf.get("convrot_groupsize", params_conf.get("convrot_groupsize", 256))
|
||||
),
|
||||
"quant_group_size": 64,
|
||||
"linear_dtype": layer_conf.get("linear_dtype", params_conf.get("linear_dtype", "int4")),
|
||||
}
|
||||
else:
|
||||
raise ValueError(f"Unsupported quantization format: {module.quant_format}")
|
||||
|
||||
@@ -1165,11 +1150,6 @@ def _quantized_weight_state_dict(module, sd, prefix, extra_quant_conf=None, extr
|
||||
if module.quant_format == "int8_tensorwise" and getattr(params, "convrot", False):
|
||||
quant_conf["convrot"] = True
|
||||
quant_conf["convrot_groupsize"] = getattr(params, "convrot_groupsize", 256)
|
||||
elif module.quant_format == "convrot_w4a4":
|
||||
quant_conf["convrot_groupsize"] = getattr(params, "convrot_groupsize", 256)
|
||||
linear_dtype = getattr(params, "linear_dtype", "int4")
|
||||
if linear_dtype != "int4":
|
||||
quant_conf["linear_dtype"] = linear_dtype
|
||||
if extra_quant_conf:
|
||||
quant_conf.update(extra_quant_conf)
|
||||
sd[f"{prefix}comfy_quant"] = torch.tensor(list(json.dumps(quant_conf).encode("utf-8")), dtype=torch.uint8)
|
||||
@@ -1257,7 +1237,7 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
|
||||
run_every_op()
|
||||
|
||||
input_shape = input.shape
|
||||
reshaped_nd = False
|
||||
reshaped_3d = False
|
||||
#If cast needs to apply lora, it should be done in the compute dtype
|
||||
compute_dtype = input.dtype
|
||||
|
||||
@@ -1294,12 +1274,12 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
|
||||
# Inference path (unchanged)
|
||||
if _use_quantized and quantize_input:
|
||||
|
||||
# Reshape >=3D tensors to 2D for quantization (needed for NVFP4 and others)
|
||||
input_reshaped = input.reshape(-1, input_shape[-1]) if input.ndim >= 3 else input
|
||||
# Reshape 3D tensors to 2D for quantization (needed for NVFP4 and others)
|
||||
input_reshaped = input.reshape(-1, input_shape[2]) if input.ndim == 3 else input
|
||||
|
||||
# Fall back to non-quantized for non-2D tensors
|
||||
if input_reshaped.ndim == 2:
|
||||
reshaped_nd = input.ndim >= 3
|
||||
reshaped_3d = input.ndim == 3
|
||||
# dtype is now implicit in the layout class
|
||||
scale = getattr(self, 'input_scale', None)
|
||||
if scale is not None:
|
||||
@@ -1314,9 +1294,9 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
|
||||
weight_only_quant=weight_only_quant,
|
||||
)
|
||||
|
||||
# Reshape output back to original rank if input was >2D
|
||||
if reshaped_nd:
|
||||
output = output.reshape((*input_shape[:-1], self.weight.shape[0]))
|
||||
# Reshape output back to 3D if input was 3D
|
||||
if reshaped_3d:
|
||||
output = output.reshape((input_shape[0], input_shape[1], self.weight.shape[0]))
|
||||
|
||||
return output
|
||||
|
||||
@@ -1450,12 +1430,6 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
|
||||
}
|
||||
if hasattr(params, "block_scale"): # NVFP4
|
||||
kwargs["block_scale"] = params.block_scale[i]
|
||||
if hasattr(params, "quant_group_size"):
|
||||
kwargs["quant_group_size"] = params.quant_group_size
|
||||
if hasattr(params, "convrot_groupsize"):
|
||||
kwargs["convrot_groupsize"] = params.convrot_groupsize
|
||||
if hasattr(params, "linear_dtype"):
|
||||
kwargs["linear_dtype"] = params.linear_dtype
|
||||
return QuantizedTensor(weight._qdata[i], weight._layout_cls, type(params)(**kwargs))
|
||||
|
||||
def state_dict(self, *args, destination=None, prefix="", **kwargs):
|
||||
|
||||
+2
-44
@@ -3,22 +3,6 @@ import logging
|
||||
|
||||
from comfy.cli_args import args
|
||||
|
||||
|
||||
def _rocm_kitchen_arch_supported():
|
||||
"""comfy-kitchen's INT8 Triton kernels compile tl.dot to matrix-core instructions.
|
||||
RDNA3/3.5/4 (gfx11xx/gfx12xx) have WMMA and CDNA (gfx9xx) has MFMA; RDNA1/RDNA2
|
||||
(gfx10xx) have neither, so the INT8 path hangs the GPU there. Gates the automatic
|
||||
ROCm default so those cards stay on the eager fallback (an explicit
|
||||
--enable-triton-backend still forces it on any arch)."""
|
||||
try:
|
||||
arch = torch.cuda.get_device_properties(torch.cuda.current_device()).gcnArchName.split(":")[0]
|
||||
except Exception:
|
||||
return False
|
||||
if arch.startswith(("gfx11", "gfx12")):
|
||||
return True
|
||||
return arch in ("gfx908", "gfx90a", "gfx940", "gfx941", "gfx942", "gfx950")
|
||||
|
||||
|
||||
try:
|
||||
import comfy_kitchen as ck
|
||||
from comfy_kitchen.tensor import (
|
||||
@@ -26,7 +10,6 @@ try:
|
||||
QuantizedLayout,
|
||||
TensorCoreFP8Layout as _CKFp8Layout,
|
||||
TensorCoreNVFP4Layout as _CKNvfp4Layout,
|
||||
TensorCoreConvRotW4A4Layout as _CKTensorCoreConvRotW4A4Layout,
|
||||
TensorWiseINT8Layout as _CKTensorWiseINT8Layout,
|
||||
register_layout_op,
|
||||
register_layout_class,
|
||||
@@ -41,22 +24,10 @@ try:
|
||||
ck.registry.disable("cuda")
|
||||
logging.warning("WARNING: You need pytorch with cu130 or higher to use optimized CUDA operations.")
|
||||
|
||||
# On ROCm/AMD the CUDA backend is unavailable, so Triton is the only accelerated
|
||||
# comfy-kitchen backend. Enable it by default there, but only on Triton >= 3.7 AND a
|
||||
# matrix-core GPU (RDNA3+ WMMA gfx11xx/gfx12xx, CDNA MFMA gfx9xx). RDNA1/RDNA2
|
||||
# (gfx10xx) have no WMMA -> the INT8 tl.dot path hangs the GPU, so they stay eager.
|
||||
# older Triton lacks libdevice.rint on the HIP backend and hard-crashes the INT8 path.
|
||||
if args.disable_triton_backend:
|
||||
ck.registry.disable("triton")
|
||||
elif args.enable_triton_backend: # or (torch.version.hip is not None and _rocm_kitchen_arch_supported()):
|
||||
if args.enable_triton_backend:
|
||||
try:
|
||||
import triton
|
||||
triton_version = tuple(int(v) for v in triton.__version__.split(".")[:2])
|
||||
if args.enable_triton_backend or triton_version >= (3, 7):
|
||||
logging.info("Found triton %s. Enabling comfy-kitchen triton backend.", triton.__version__)
|
||||
else:
|
||||
logging.info("Triton %s is too old for the ROCm INT8 path (needs >= 3.7); comfy-kitchen triton backend disabled.", triton.__version__)
|
||||
ck.registry.disable("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")
|
||||
@@ -80,9 +51,6 @@ except ImportError as e:
|
||||
class _CKTensorWiseINT8Layout:
|
||||
pass
|
||||
|
||||
class _CKTensorCoreConvRotW4A4Layout:
|
||||
pass
|
||||
|
||||
def register_layout_class(name, cls):
|
||||
pass
|
||||
|
||||
@@ -211,7 +179,6 @@ class TensorCoreFP8E5M2Layout(_TensorCoreFP8LayoutBase):
|
||||
# Backward compatibility alias - default to E4M3
|
||||
TensorCoreFP8Layout = TensorCoreFP8E4M3Layout
|
||||
TensorWiseINT8Layout = _CKTensorWiseINT8Layout
|
||||
TensorCoreConvRotW4A4Layout = _CKTensorCoreConvRotW4A4Layout
|
||||
|
||||
|
||||
# ==============================================================================
|
||||
@@ -223,7 +190,6 @@ register_layout_class("TensorCoreFP8E4M3Layout", TensorCoreFP8E4M3Layout)
|
||||
register_layout_class("TensorCoreFP8E5M2Layout", TensorCoreFP8E5M2Layout)
|
||||
register_layout_class("TensorCoreNVFP4Layout", TensorCoreNVFP4Layout)
|
||||
register_layout_class("TensorWiseINT8Layout", _CKTensorWiseINT8Layout)
|
||||
register_layout_class("TensorCoreConvRotW4A4Layout", _CKTensorCoreConvRotW4A4Layout)
|
||||
if _CK_MXFP8_AVAILABLE:
|
||||
register_layout_class("TensorCoreMXFP8Layout", TensorCoreMXFP8Layout)
|
||||
|
||||
@@ -261,13 +227,6 @@ QUANT_ALGOS["int8_tensorwise"] = {
|
||||
"quantize_input": False,
|
||||
}
|
||||
|
||||
QUANT_ALGOS["convrot_w4a4"] = {
|
||||
"storage_t": torch.int8,
|
||||
"parameters": {"weight_scale"},
|
||||
"comfy_tensor_layout": "TensorCoreConvRotW4A4Layout",
|
||||
"quantize_input": False,
|
||||
}
|
||||
|
||||
|
||||
# ==============================================================================
|
||||
# Re-exports for backward compatibility
|
||||
@@ -280,7 +239,6 @@ __all__ = [
|
||||
"TensorCoreFP8E4M3Layout",
|
||||
"TensorCoreFP8E5M2Layout",
|
||||
"TensorCoreNVFP4Layout",
|
||||
"TensorCoreConvRotW4A4Layout",
|
||||
"TensorWiseINT8Layout",
|
||||
"QUANT_ALGOS",
|
||||
"register_layout_op",
|
||||
|
||||
+19
-83
@@ -16,7 +16,6 @@ import comfy.ldm.cosmos.vae
|
||||
import comfy.ldm.wan.vae
|
||||
import comfy.ldm.wan.vae2_2
|
||||
import comfy.ldm.hunyuan3d.vae
|
||||
import comfy.ldm.seedvr.vae
|
||||
import comfy.ldm.triposplat.vae
|
||||
import comfy.ldm.ace.vae.music_dcae_pipeline
|
||||
import comfy.ldm.cogvideo.vae
|
||||
@@ -474,8 +473,7 @@ class CLIP:
|
||||
|
||||
class VAE:
|
||||
def __init__(self, sd=None, device=None, config=None, dtype=None, metadata=None):
|
||||
is_seedvr2_vae = "decoder.up_blocks.2.upsamplers.0.upscale_conv.weight" in sd
|
||||
if not is_seedvr2_vae and 'decoder.up_blocks.0.resnets.0.norm1.weight' in sd.keys(): #diffusers format
|
||||
if 'decoder.up_blocks.0.resnets.0.norm1.weight' in sd.keys(): #diffusers format
|
||||
sd = diffusers_convert.convert_vae_state_dict(sd)
|
||||
|
||||
if model_management.is_amd():
|
||||
@@ -502,8 +500,6 @@ class VAE:
|
||||
self.upscale_index_formula = None
|
||||
self.extra_1d_channel = None
|
||||
self.crop_input = True
|
||||
self.handles_tiling = False
|
||||
self.format_encoded = None
|
||||
|
||||
self.audio_sample_rate = 44100
|
||||
|
||||
@@ -550,22 +546,6 @@ class VAE:
|
||||
self.first_stage_model = StageC_coder()
|
||||
self.downscale_ratio = 32
|
||||
self.latent_channels = 16
|
||||
elif "decoder.up_blocks.2.upsamplers.0.upscale_conv.weight" in sd: # seedvr2
|
||||
self.first_stage_model = comfy.ldm.seedvr.vae.VideoAutoencoderKLWrapper()
|
||||
self.latent_channels = comfy.ldm.seedvr.vae.SEEDVR2_LATENT_CHANNELS
|
||||
self.latent_dim = 3
|
||||
self.disable_offload = True
|
||||
self.memory_used_decode = lambda shape, dtype: self.first_stage_model.comfy_memory_used_decode(shape)
|
||||
self.memory_used_encode = lambda shape, dtype: (max(shape[2], 5) * shape[3] * shape[4] * 64) * model_management.dtype_size(dtype)
|
||||
self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32]
|
||||
self.handles_tiling = True
|
||||
self.format_encoded = self.first_stage_model.comfy_format_encoded
|
||||
self.downscale_ratio = (lambda a: max(0, math.floor((a + 3) / 4)), 8, 8)
|
||||
self.downscale_index_formula = (4, 8, 8)
|
||||
self.upscale_ratio = (lambda a: max(0, a * 4 - 3), 8, 8)
|
||||
self.upscale_index_formula = (4, 8, 8)
|
||||
self.process_input = lambda image: image * 2.0 - 1.0
|
||||
self.crop_input = False
|
||||
elif "decoder.conv_in.weight" in sd:
|
||||
if sd['decoder.conv_in.weight'].shape[1] == 64:
|
||||
ddconfig = {"block_out_channels": [128, 256, 512, 512, 1024, 1024], "in_channels": 3, "out_channels": 3, "num_res_blocks": 2, "ffactor_spatial": 32, "downsample_match_channel": True, "upsample_match_channel": True}
|
||||
@@ -1032,10 +1012,6 @@ class VAE:
|
||||
decode_fn = lambda a: self.first_stage_model.decode(a.to(self.vae_dtype).to(self.device)).to(dtype=self.vae_output_dtype())
|
||||
return self.process_output(comfy.utils.tiled_scale_multidim(samples, decode_fn, tile=(tile_t, tile_x, tile_y), overlap=overlap, upscale_amount=self.upscale_ratio, out_channels=self.output_channels, index_formulas=self.upscale_index_formula, output_device=self.output_device))
|
||||
|
||||
def _decode_tiled_owned(self, samples, **kwargs):
|
||||
out = self.first_stage_model.decode_tiled(samples.to(self.vae_dtype).to(self.device), **kwargs)
|
||||
return self.process_output(out.to(device=self.output_device, dtype=self.vae_output_dtype(), copy=True))
|
||||
|
||||
def encode_tiled_(self, pixel_samples, tile_x=512, tile_y=512, overlap = 64):
|
||||
steps = pixel_samples.shape[0] * comfy.utils.get_tiled_scale_steps(pixel_samples.shape[3], pixel_samples.shape[2], tile_x, tile_y, overlap)
|
||||
steps += pixel_samples.shape[0] * comfy.utils.get_tiled_scale_steps(pixel_samples.shape[3], pixel_samples.shape[2], tile_x // 2, tile_y * 2, overlap)
|
||||
@@ -1072,25 +1048,6 @@ class VAE:
|
||||
encode_fn = lambda a: self.first_stage_model.encode((self.process_input(a)).to(self.vae_dtype).to(self.device)).to(dtype=self.vae_output_dtype())
|
||||
return comfy.utils.tiled_scale_multidim(samples, encode_fn, tile=(tile_t, tile_x, tile_y), overlap=overlap, upscale_amount=self.downscale_ratio, out_channels=self.latent_channels, downscale=True, index_formulas=self.downscale_index_formula, output_device=self.output_device)
|
||||
|
||||
def _encode_tiled_owned(self, pixel_samples, **kwargs):
|
||||
x = self.process_input(pixel_samples).to(self.vae_dtype).to(self.device)
|
||||
out = self.first_stage_model.encode_tiled(x, **kwargs)
|
||||
return out.to(device=self.output_device, dtype=self.vae_output_dtype())
|
||||
|
||||
def _owned_tiled_args(self, tile_x=None, tile_y=None, overlap=None, tile_t=None, overlap_t=None):
|
||||
args = {}
|
||||
if tile_x is not None:
|
||||
args["tile_x"] = tile_x
|
||||
if tile_y is not None:
|
||||
args["tile_y"] = tile_y
|
||||
if overlap is not None:
|
||||
args["overlap"] = overlap
|
||||
if tile_t is not None:
|
||||
args["tile_t"] = tile_t
|
||||
if overlap_t is not None:
|
||||
args["overlap_t"] = overlap_t
|
||||
return args
|
||||
|
||||
def decode(self, samples_in, vae_options={}):
|
||||
self.throw_exception_if_invalid()
|
||||
pixel_samples = None
|
||||
@@ -1138,19 +1095,11 @@ class VAE:
|
||||
if dims == 1 or self.extra_1d_channel is not None:
|
||||
pixel_samples = self.decode_tiled_1d(samples_in)
|
||||
elif dims == 2:
|
||||
if self.handles_tiling:
|
||||
tile = 256 // self.spacial_compression_decode()
|
||||
overlap = tile // 4
|
||||
pixel_samples = self._decode_tiled_owned(samples_in, tile_x=tile, tile_y=tile, overlap=overlap)
|
||||
else:
|
||||
pixel_samples = self.decode_tiled_(samples_in)
|
||||
pixel_samples = self.decode_tiled_(samples_in)
|
||||
elif dims == 3:
|
||||
tile = 256 // self.spacial_compression_decode()
|
||||
overlap = tile // 4
|
||||
if self.handles_tiling:
|
||||
pixel_samples = self._decode_tiled_owned(samples_in, tile_x=tile, tile_y=tile, overlap=overlap)
|
||||
else:
|
||||
pixel_samples = self.decode_tiled_3d(samples_in, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap))
|
||||
pixel_samples = self.decode_tiled_3d(samples_in, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap))
|
||||
|
||||
pixel_samples = pixel_samples.to(self.output_device).movedim(1,-1)
|
||||
return pixel_samples
|
||||
@@ -1169,9 +1118,7 @@ class VAE:
|
||||
args["overlap"] = overlap
|
||||
|
||||
with model_management.cuda_device_context(self.device):
|
||||
if self.handles_tiling and dims in (2, 3):
|
||||
output = self._decode_tiled_owned(samples, **self._owned_tiled_args(tile_x, tile_y, overlap, tile_t, overlap_t))
|
||||
elif dims == 1 or self.extra_1d_channel is not None:
|
||||
if dims == 1 or self.extra_1d_channel is not None:
|
||||
args.pop("tile_y")
|
||||
output = self.decode_tiled_1d(samples, **args)
|
||||
elif dims == 2:
|
||||
@@ -1232,17 +1179,12 @@ class VAE:
|
||||
if self.latent_dim == 3:
|
||||
tile = 256
|
||||
overlap = tile // 4
|
||||
if self.handles_tiling:
|
||||
samples = self._encode_tiled_owned(pixel_samples, tile_x=tile, tile_y=tile, overlap=overlap)
|
||||
else:
|
||||
samples = self.encode_tiled_3d(pixel_samples, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap))
|
||||
samples = self.encode_tiled_3d(pixel_samples, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap))
|
||||
elif self.latent_dim == 1 or self.extra_1d_channel is not None:
|
||||
samples = self.encode_tiled_1d(pixel_samples)
|
||||
else:
|
||||
samples = self.encode_tiled_(pixel_samples)
|
||||
|
||||
if self.format_encoded is not None:
|
||||
samples = self.format_encoded(samples)
|
||||
return samples
|
||||
|
||||
def encode_tiled(self, pixel_samples, tile_x=None, tile_y=None, overlap=None, tile_t=None, overlap_t=None):
|
||||
@@ -1250,7 +1192,7 @@ class VAE:
|
||||
pixel_samples = self.vae_encode_crop_pixels(pixel_samples)
|
||||
dims = self.latent_dim
|
||||
pixel_samples = pixel_samples.movedim(-1, 1)
|
||||
if dims == 3 and pixel_samples.ndim < 5:
|
||||
if dims == 3:
|
||||
if not self.not_video:
|
||||
pixel_samples = pixel_samples.movedim(1, 0).unsqueeze(0)
|
||||
else:
|
||||
@@ -1274,27 +1216,21 @@ class VAE:
|
||||
elif dims == 2:
|
||||
samples = self.encode_tiled_(pixel_samples, **args)
|
||||
elif dims == 3:
|
||||
if self.handles_tiling:
|
||||
samples = self._encode_tiled_owned(pixel_samples, **self._owned_tiled_args(tile_x, tile_y, overlap, tile_t, overlap_t))
|
||||
if tile_t is not None:
|
||||
tile_t_latent = max(2, self.downscale_ratio[0](tile_t))
|
||||
else:
|
||||
if tile_t is not None:
|
||||
tile_t_latent = max(2, self.downscale_ratio[0](tile_t))
|
||||
else:
|
||||
tile_t_latent = 9999
|
||||
args["tile_t"] = self.upscale_ratio[0](tile_t_latent)
|
||||
tile_t_latent = 9999
|
||||
args["tile_t"] = self.upscale_ratio[0](tile_t_latent)
|
||||
|
||||
spatial_overlap = overlap if overlap is not None else 64
|
||||
if overlap_t is None:
|
||||
args["overlap"] = (1, spatial_overlap, spatial_overlap)
|
||||
else:
|
||||
args["overlap"] = (self.upscale_ratio[0](max(1, min(tile_t_latent // 2, self.downscale_ratio[0](overlap_t)))), spatial_overlap, spatial_overlap)
|
||||
maximum = pixel_samples.shape[2]
|
||||
maximum = self.upscale_ratio[0](self.downscale_ratio[0](maximum))
|
||||
if overlap_t is None:
|
||||
args["overlap"] = (1, overlap, overlap)
|
||||
else:
|
||||
args["overlap"] = (self.upscale_ratio[0](max(1, min(tile_t_latent // 2, self.downscale_ratio[0](overlap_t)))), overlap, overlap)
|
||||
maximum = pixel_samples.shape[2]
|
||||
maximum = self.upscale_ratio[0](self.downscale_ratio[0](maximum))
|
||||
|
||||
samples = self.encode_tiled_3d(pixel_samples[:,:,:maximum], **args)
|
||||
samples = self.encode_tiled_3d(pixel_samples[:,:,:maximum], **args)
|
||||
|
||||
if self.format_encoded is not None:
|
||||
samples = self.format_encoded(samples)
|
||||
return samples
|
||||
|
||||
def get_sd(self):
|
||||
@@ -1962,7 +1898,7 @@ def load_state_dict_guess_config(sd, output_vae=True, output_clip=True, output_c
|
||||
manual_cast_dtype = model_management.unet_manual_cast(None, load_device, model_config.supported_inference_dtypes)
|
||||
else:
|
||||
manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)
|
||||
model_config.set_inference_dtype(unet_dtype, manual_cast_dtype, device=load_device)
|
||||
model_config.set_inference_dtype(unet_dtype, manual_cast_dtype)
|
||||
|
||||
if model_config.clip_vision_prefix is not None:
|
||||
if output_clipvision:
|
||||
@@ -2103,7 +2039,7 @@ def load_diffusion_model_state_dict(sd, model_options={}, metadata=None, disable
|
||||
manual_cast_dtype = model_management.unet_manual_cast(None, load_device, model_config.supported_inference_dtypes)
|
||||
else:
|
||||
manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)
|
||||
model_config.set_inference_dtype(unet_dtype, manual_cast_dtype, device=load_device)
|
||||
model_config.set_inference_dtype(unet_dtype, manual_cast_dtype)
|
||||
|
||||
if custom_operations is not None:
|
||||
model_config.custom_operations = custom_operations
|
||||
|
||||
@@ -1685,40 +1685,6 @@ class Chroma(supported_models_base.BASE):
|
||||
t5_detect = comfy.text_encoders.sd3_clip.t5_xxl_detect(state_dict, "{}t5xxl.transformer.".format(pref))
|
||||
return supported_models_base.ClipTarget(comfy.text_encoders.pixart_t5.PixArtTokenizer, comfy.text_encoders.pixart_t5.pixart_te(**t5_detect))
|
||||
|
||||
class SeedVR2(supported_models_base.BASE):
|
||||
unet_config = {
|
||||
"image_model": "seedvr2"
|
||||
}
|
||||
unet_extra_config = {}
|
||||
required_keys = {
|
||||
"{}positive_conditioning",
|
||||
"{}negative_conditioning",
|
||||
}
|
||||
latent_format = comfy.latent_formats.SeedVR2
|
||||
|
||||
vae_key_prefix = ["vae."]
|
||||
text_encoder_key_prefix = ["text_encoders."]
|
||||
supported_inference_dtypes = [torch.bfloat16, torch.float16, torch.float32]
|
||||
sampling_settings = {
|
||||
"shift": 1.0,
|
||||
}
|
||||
|
||||
def set_inference_dtype(self, dtype, manual_cast_dtype, device=None):
|
||||
if (
|
||||
dtype == torch.float16
|
||||
and manual_cast_dtype is None
|
||||
and comfy.model_management.should_use_bf16(device)
|
||||
):
|
||||
manual_cast_dtype = torch.bfloat16
|
||||
super().set_inference_dtype(dtype, manual_cast_dtype, device=device)
|
||||
|
||||
def get_model(self, state_dict, prefix="", device=None):
|
||||
out = model_base.SeedVR2(self, device=device)
|
||||
return out
|
||||
|
||||
def clip_target(self, state_dict={}):
|
||||
return None
|
||||
|
||||
class ChromaRadiance(Chroma):
|
||||
unet_config = {
|
||||
"image_model": "chroma_radiance",
|
||||
@@ -2382,7 +2348,6 @@ models = [
|
||||
HiDream,
|
||||
HiDreamO1,
|
||||
Chroma,
|
||||
SeedVR2,
|
||||
ChromaRadiance,
|
||||
ACEStep,
|
||||
ACEStep15,
|
||||
|
||||
@@ -54,13 +54,13 @@ class BASE:
|
||||
optimizations = {"fp8": False}
|
||||
|
||||
@classmethod
|
||||
def matches(s, unet_config, state_dict=None, unet_key_prefix=""):
|
||||
def matches(s, unet_config, state_dict=None):
|
||||
for k in s.unet_config:
|
||||
if k not in unet_config or s.unet_config[k] != unet_config[k]:
|
||||
return False
|
||||
if state_dict is not None:
|
||||
for k in s.required_keys:
|
||||
if k.format(unet_key_prefix) not in state_dict:
|
||||
if k not in state_dict:
|
||||
return False
|
||||
return True
|
||||
|
||||
@@ -115,7 +115,7 @@ class BASE:
|
||||
replace_prefix = {"": self.vae_key_prefix[0]}
|
||||
return utils.state_dict_prefix_replace(state_dict, replace_prefix)
|
||||
|
||||
def set_inference_dtype(self, dtype, manual_cast_dtype, device=None):
|
||||
def set_inference_dtype(self, dtype, manual_cast_dtype):
|
||||
self.unet_config['dtype'] = dtype
|
||||
self.manual_cast_dtype = manual_cast_dtype
|
||||
|
||||
|
||||
@@ -1088,7 +1088,7 @@ class Gemma4_Tokenizer():
|
||||
h, w = samples.shape[2], samples.shape[3]
|
||||
patch_size = 16
|
||||
pooling_k = 3
|
||||
max_soft_tokens = kwargs.get("max_soft_tokens", 70 if is_video else 280)
|
||||
max_soft_tokens = 70 if is_video else 280 # video uses smaller token budget per frame
|
||||
max_patches = max_soft_tokens * pooling_k * pooling_k
|
||||
target_px = max_patches * patch_size * patch_size
|
||||
factor = (target_px / (h * w)) ** 0.5
|
||||
|
||||
@@ -90,27 +90,6 @@ class Qwen3VL(BaseLlama, BaseQwen3, BaseGenerate, torch.nn.Module):
|
||||
deepstack = [torch.cat([deepstack[i], ds[i]], dim=0) for i in range(len(ds))]
|
||||
return position_ids, visual_pos_masks, deepstack
|
||||
|
||||
def forward(self, input_ids, attention_mask=None, embeds=None, num_tokens=None, intermediate_output=None, final_layer_norm_intermediate=True, dtype=None, embeds_info=[], **kwargs):
|
||||
position_ids = kwargs.pop("position_ids", None)
|
||||
visual_pos_masks = kwargs.pop("visual_pos_masks", None)
|
||||
deepstack_embeds = kwargs.pop("deepstack_embeds", None)
|
||||
if embeds is not None and position_ids is None:
|
||||
position_ids, visual_pos_masks, deepstack_embeds = self.build_image_inputs(embeds, embeds_info)
|
||||
return self.model(
|
||||
input_ids,
|
||||
attention_mask=attention_mask,
|
||||
embeds=embeds,
|
||||
num_tokens=num_tokens,
|
||||
intermediate_output=intermediate_output,
|
||||
final_layer_norm_intermediate=final_layer_norm_intermediate,
|
||||
dtype=dtype,
|
||||
position_ids=position_ids,
|
||||
embeds_info=embeds_info,
|
||||
visual_pos_masks=visual_pos_masks,
|
||||
deepstack_embeds=deepstack_embeds,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
def _make_qwen3vl_model(model_type):
|
||||
class Qwen3VL_(Qwen3VL):
|
||||
|
||||
@@ -100,7 +100,6 @@ def _parse_cli_feature_flags() -> dict[str, Any]:
|
||||
# Default server capabilities
|
||||
_CORE_FEATURE_FLAGS: dict[str, Any] = {
|
||||
"supports_preview_metadata": True,
|
||||
"supports_node_failure_policy": True,
|
||||
"supports_model_type_tags": True,
|
||||
"max_upload_size": args.max_upload_size * 1024 * 1024, # Convert MB to bytes
|
||||
"extension": {"manager": {"supports_v4": True}},
|
||||
|
||||
@@ -17,10 +17,6 @@ class Seedream4Options(BaseModel):
|
||||
max_images: int = Field(15)
|
||||
|
||||
|
||||
class Seedream5OptimizePromptOptions(BaseModel):
|
||||
thinking: Literal["auto", "enabled", "disabled"] = Field(...)
|
||||
|
||||
|
||||
class Seedream4TaskCreationRequest(BaseModel):
|
||||
model: str = Field(...)
|
||||
prompt: str = Field(...)
|
||||
@@ -32,7 +28,6 @@ class Seedream4TaskCreationRequest(BaseModel):
|
||||
sequential_image_generation_options: Seedream4Options | None = Field(Seedream4Options(max_images=15))
|
||||
watermark: bool = Field(False)
|
||||
output_format: str | None = None
|
||||
optimize_prompt_options: Seedream5OptimizePromptOptions | None = None
|
||||
|
||||
|
||||
class ImageTaskCreationResponse(BaseModel):
|
||||
|
||||
@@ -77,7 +77,6 @@ class To3DUVTaskRequest(BaseModel):
|
||||
|
||||
class To3DPartTaskRequest(BaseModel):
|
||||
File: TaskFile3DInput = Field(...)
|
||||
EnableStagedGeneration: bool | None = Field(None)
|
||||
|
||||
|
||||
class TextureEditImageInfo(BaseModel):
|
||||
|
||||
@@ -34,7 +34,6 @@ from comfy_api_nodes.apis.bytedance import (
|
||||
SeedanceVirtualLibraryCreateAssetRequest,
|
||||
Seedream4Options,
|
||||
Seedream4TaskCreationRequest,
|
||||
Seedream5OptimizePromptOptions,
|
||||
TaskAudioContent,
|
||||
TaskAudioContentUrl,
|
||||
TaskCreationResponse,
|
||||
@@ -876,17 +875,6 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode):
|
||||
tooltip='Whether to add an "AI generated" watermark to the image.',
|
||||
advanced=True,
|
||||
),
|
||||
IO.Boolean.Input(
|
||||
"thinking",
|
||||
default=True,
|
||||
tooltip=(
|
||||
"Enable the model's prompt-optimization reasoning ('thinking') for better adherence. "
|
||||
"Can substantially increase generation time — notably on Seedream 5.0 Pro. "
|
||||
"Can only be disabled for text-to-image (not when reference images are provided)."
|
||||
),
|
||||
optional=True,
|
||||
advanced=True,
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
IO.Image.Output(),
|
||||
@@ -932,7 +920,6 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode):
|
||||
model: dict,
|
||||
seed: int = 0,
|
||||
watermark: bool = False,
|
||||
thinking: bool = True,
|
||||
) -> IO.NodeOutput:
|
||||
validate_string(prompt, strip_whitespace=True, min_length=1)
|
||||
model_id = SEEDREAM_MODELS[model["model"]]
|
||||
@@ -992,10 +979,6 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode):
|
||||
raise ValueError(
|
||||
"The maximum number of generated images plus the number of reference images cannot exceed 15."
|
||||
)
|
||||
if not thinking and n_input_images > 0:
|
||||
raise ValueError(
|
||||
"'thinking' can only be disabled for text-to-image; enable it when using reference images."
|
||||
)
|
||||
|
||||
reference_images_urls: list[str] = []
|
||||
if image_tensors:
|
||||
@@ -1009,9 +992,6 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode):
|
||||
wait_label="Uploading reference images",
|
||||
)
|
||||
|
||||
optimize_prompt_options = None
|
||||
if n_input_images == 0:
|
||||
optimize_prompt_options = Seedream5OptimizePromptOptions(thinking="enabled" if thinking else "disabled")
|
||||
response = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(path=BYTEPLUS_IMAGE_ENDPOINT, method="POST"),
|
||||
@@ -1025,7 +1005,6 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode):
|
||||
sequential_image_generation=None if is_pro else sequential_image_generation,
|
||||
sequential_image_generation_options=None if is_pro else Seedream4Options(max_images=max_images),
|
||||
watermark=watermark,
|
||||
optimize_prompt_options=optimize_prompt_options,
|
||||
),
|
||||
)
|
||||
if len(response.data) == 1:
|
||||
|
||||
@@ -1133,9 +1133,7 @@ class GeminiImage2(IO.ComfyNode):
|
||||
) -> IO.NodeOutput:
|
||||
validate_string(prompt, strip_whitespace=True, min_length=1)
|
||||
if model == "Nano Banana 2 (Gemini 3.1 Flash Image)":
|
||||
model = "gemini-3.1-flash-image"
|
||||
elif model == "gemini-3-pro-image-preview":
|
||||
model = "gemini-3-pro-image"
|
||||
model = "gemini-3.1-flash-image-preview"
|
||||
|
||||
parts: list[GeminiPart] = [GeminiPart(text=prompt)]
|
||||
if images is not None:
|
||||
@@ -1509,7 +1507,7 @@ class GeminiNanoBanana2V2(IO.ComfyNode):
|
||||
validate_string(prompt, strip_whitespace=True, min_length=1)
|
||||
model_choice = model["model"]
|
||||
if model_choice == "Nano Banana 2 (Gemini 3.1 Flash Image)":
|
||||
model_id = "gemini-3.1-flash-image"
|
||||
model_id = "gemini-3.1-flash-image-preview"
|
||||
elif model_choice == "Nano Banana 2 Lite":
|
||||
model_id = "gemini-3.1-flash-lite-image"
|
||||
else:
|
||||
|
||||
@@ -642,7 +642,6 @@ class Tencent3DPartNode(IO.ComfyNode):
|
||||
response_model=To3DProTaskCreateResponse,
|
||||
data=To3DPartTaskRequest(
|
||||
File=TaskFile3DInput(Type=file_format.upper(), Url=model_url),
|
||||
EnableStagedGeneration=True,
|
||||
),
|
||||
is_rate_limited=_is_tencent_rate_limited,
|
||||
)
|
||||
|
||||
@@ -11,11 +11,9 @@ from io import BytesIO
|
||||
from yarl import URL
|
||||
|
||||
from comfy.cli_args import args
|
||||
from comfy.comfy_api_env import normalize_comfy_api_base
|
||||
from comfy.deploy_environment import get_deploy_environment
|
||||
from comfy.model_management import processing_interrupted
|
||||
from comfy_api.latest import IO
|
||||
from comfyui_version import __version__ as comfyui_version
|
||||
|
||||
from .common_exceptions import ProcessingInterrupted
|
||||
|
||||
@@ -61,12 +59,11 @@ def get_comfy_api_headers(node_cls: type[IO.ComfyNode]) -> dict[str, str]:
|
||||
**get_auth_header(node_cls),
|
||||
"Comfy-Env": get_deploy_environment(),
|
||||
"Comfy-Usage-Source": get_usage_source(node_cls),
|
||||
"Comfy-Core-Version": comfyui_version,
|
||||
}
|
||||
|
||||
|
||||
def default_base_url() -> str:
|
||||
return normalize_comfy_api_base(getattr(args, "comfy_api_base", "https://api.comfy.org"))
|
||||
return getattr(args, "comfy_api_base", "https://api.comfy.org")
|
||||
|
||||
|
||||
async def sleep_with_interrupt(
|
||||
|
||||
@@ -503,8 +503,6 @@ RAM_CACHE_DEFAULT_RAM_USAGE = 0.05
|
||||
|
||||
RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER = 1.3
|
||||
|
||||
RAM_CACHE_LARGE_INTERMEDIATE = 512 * 1024 ** 2
|
||||
|
||||
|
||||
def all_outputs_dynamic(outputs):
|
||||
if outputs is None:
|
||||
@@ -519,6 +517,7 @@ def all_outputs_dynamic(outputs):
|
||||
|
||||
return True
|
||||
|
||||
|
||||
class RAMPressureCache(LRUCache):
|
||||
|
||||
def __init__(self, key_class, enable_providers=False):
|
||||
@@ -540,9 +539,9 @@ class RAMPressureCache(LRUCache):
|
||||
self.timestamps[self.cache_key_set.get_data_key(node_id)] = time.time()
|
||||
super().set_local(node_id, value)
|
||||
|
||||
def ram_release(self, target, free_active=False, min_entry_size=0):
|
||||
def ram_release(self, target, free_active=False):
|
||||
if psutil.virtual_memory().available >= target:
|
||||
return 0
|
||||
return
|
||||
|
||||
clean_list = []
|
||||
|
||||
@@ -556,9 +555,8 @@ class RAMPressureCache(LRUCache):
|
||||
oom_score = RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER ** (self.generation - self.used_generation[key])
|
||||
|
||||
ram_usage = RAM_CACHE_DEFAULT_RAM_USAGE
|
||||
oom_ram_usage = ram_usage
|
||||
def scan_list_for_ram_usage(outputs):
|
||||
nonlocal ram_usage, oom_ram_usage
|
||||
nonlocal ram_usage
|
||||
if outputs is None:
|
||||
return
|
||||
for output in outputs:
|
||||
@@ -566,26 +564,19 @@ class RAMPressureCache(LRUCache):
|
||||
scan_list_for_ram_usage(output)
|
||||
elif isinstance(output, torch.Tensor) and output.device.type == 'cpu':
|
||||
ram_usage += output.numel() * output.element_size()
|
||||
oom_ram_usage += output.numel() * output.element_size()
|
||||
elif isinstance(output, ModelPatcher) and self.used_generation[key] != self.generation:
|
||||
#old ModelPatchers are the first to go
|
||||
oom_ram_usage = 1e30
|
||||
ram_usage = 1e30
|
||||
scan_list_for_ram_usage(cache_entry.outputs)
|
||||
|
||||
if ram_usage < min_entry_size:
|
||||
continue
|
||||
|
||||
oom_score *= oom_ram_usage
|
||||
oom_score *= ram_usage
|
||||
#In the case where we have no information on the node ram usage at all,
|
||||
#break OOM score ties on the last touch timestamp (pure LRU)
|
||||
bisect.insort(clean_list, (oom_score, self.timestamps[key], key, ram_usage))
|
||||
bisect.insort(clean_list, (oom_score, self.timestamps[key], key))
|
||||
|
||||
freed = 0
|
||||
while psutil.virtual_memory().available < target and clean_list:
|
||||
_, _, key, ram_usage = clean_list.pop()
|
||||
_, _, key = clean_list.pop()
|
||||
del self.cache[key]
|
||||
self.used_generation.pop(key, None)
|
||||
self.timestamps.pop(key, None)
|
||||
self.children.pop(key, None)
|
||||
freed += ram_usage
|
||||
return freed
|
||||
|
||||
@@ -3,12 +3,11 @@ from typing import Type, Literal
|
||||
import nodes
|
||||
import asyncio
|
||||
import inspect
|
||||
from comfy_execution.graph_utils import is_link, ExecutionBlocker, ExecutionFailureBlocker
|
||||
from comfy_execution.graph_utils import is_link, ExecutionBlocker
|
||||
from comfy.comfy_types.node_typing import ComfyNodeABC, InputTypeDict, InputTypeOptions
|
||||
|
||||
# NOTE: ExecutionBlocker code got moved to graph_utils.py to prevent torch being imported too soon during unit tests
|
||||
ExecutionBlocker = ExecutionBlocker
|
||||
ExecutionFailureBlocker = ExecutionFailureBlocker
|
||||
|
||||
class DependencyCycleError(Exception):
|
||||
pass
|
||||
@@ -202,28 +201,19 @@ class ExecutionList(TopologicalSort):
|
||||
self.staged_node_id = None
|
||||
self.execution_cache = {}
|
||||
self.execution_cache_listeners = {}
|
||||
self.transient_cache = {}
|
||||
self.failure_tainted_parents = set()
|
||||
|
||||
def is_cached(self, node_id):
|
||||
return node_id in self.transient_cache or self.output_cache.get_local(node_id) is not None
|
||||
|
||||
def _get_cache_value(self, node_id):
|
||||
if node_id in self.transient_cache:
|
||||
return self.transient_cache[node_id]
|
||||
return self.output_cache.get_local(node_id)
|
||||
return self.output_cache.get_local(node_id) is not None
|
||||
|
||||
def cache_link(self, from_node_id, to_node_id):
|
||||
if to_node_id not in self.execution_cache:
|
||||
self.execution_cache[to_node_id] = {}
|
||||
self.execution_cache[to_node_id][from_node_id] = self._get_cache_value(from_node_id)
|
||||
self.execution_cache[to_node_id][from_node_id] = self.output_cache.get_local(from_node_id)
|
||||
if from_node_id not in self.execution_cache_listeners:
|
||||
self.execution_cache_listeners[from_node_id] = set()
|
||||
self.execution_cache_listeners[from_node_id].add(to_node_id)
|
||||
|
||||
def get_cache(self, from_node_id, to_node_id):
|
||||
if from_node_id in self.transient_cache:
|
||||
return self.transient_cache[from_node_id]
|
||||
if to_node_id not in self.execution_cache:
|
||||
return None
|
||||
value = self.execution_cache[to_node_id].get(from_node_id)
|
||||
@@ -233,9 +223,7 @@ class ExecutionList(TopologicalSort):
|
||||
self.output_cache.set_local(from_node_id, value)
|
||||
return value
|
||||
|
||||
def cache_update(self, node_id, value, transient=False):
|
||||
if transient:
|
||||
self.transient_cache[node_id] = value
|
||||
def cache_update(self, node_id, value):
|
||||
if node_id in self.execution_cache_listeners:
|
||||
for to_node_id in self.execution_cache_listeners[node_id]:
|
||||
if to_node_id in self.execution_cache:
|
||||
@@ -245,25 +233,6 @@ class ExecutionList(TopologicalSort):
|
||||
super().add_strong_link(from_node_id, from_socket, to_node_id)
|
||||
self.cache_link(from_node_id, to_node_id)
|
||||
|
||||
def add_completion_link(self, from_node_id, to_node_id):
|
||||
# Block to_node_id until from_node_id finishes, without consuming any of its output sockets.
|
||||
if not self.is_cached(from_node_id):
|
||||
self.add_node(from_node_id)
|
||||
if to_node_id not in self.blocking[from_node_id]:
|
||||
self.blocking[from_node_id][to_node_id] = {}
|
||||
self.blockCount[to_node_id] += 1
|
||||
|
||||
def mark_failure_tainted(self, node_id):
|
||||
# Taint all ephemeral ancestors of a failed or failure-blocked node so dynamically-expanded parents are
|
||||
# never cached as reusable when part of their expansion did not complete.
|
||||
parent_id = self.dynprompt.get_parent_node_id(node_id)
|
||||
while parent_id is not None and parent_id not in self.failure_tainted_parents:
|
||||
self.failure_tainted_parents.add(parent_id)
|
||||
parent_id = self.dynprompt.get_parent_node_id(parent_id)
|
||||
|
||||
def is_failure_tainted(self, node_id):
|
||||
return node_id in self.failure_tainted_parents
|
||||
|
||||
async def stage_node_execution(self):
|
||||
assert self.staged_node_id is None
|
||||
if self.is_empty():
|
||||
@@ -354,9 +323,7 @@ class ExecutionList(TopologicalSort):
|
||||
blocked_by = { node_id: {} for node_id in self.pendingNodes }
|
||||
for from_node_id in self.blocking:
|
||||
for to_node_id in self.blocking[from_node_id]:
|
||||
# Strong links have a True socket entry; completion links have no socket entries at all.
|
||||
sockets = self.blocking[from_node_id][to_node_id]
|
||||
if len(sockets) == 0 or True in sockets.values():
|
||||
if True in self.blocking[from_node_id][to_node_id].values():
|
||||
blocked_by[to_node_id][from_node_id] = True
|
||||
to_remove = [node_id for node_id in blocked_by if len(blocked_by[node_id]) == 0]
|
||||
while len(to_remove) > 0:
|
||||
|
||||
@@ -153,9 +153,3 @@ class ExecutionBlocker:
|
||||
"""
|
||||
def __init__(self, message):
|
||||
self.message = message
|
||||
|
||||
|
||||
class ExecutionFailureBlocker(ExecutionBlocker):
|
||||
def __init__(self, node_id):
|
||||
super().__init__(None)
|
||||
self.node_id = node_id
|
||||
|
||||
+4
-33
@@ -56,9 +56,6 @@ PREVIEWABLE_MEDIA_TYPES = frozenset({'images', 'video', 'audio', '3d', 'text'})
|
||||
# 3D file extensions for preview fallback (no dedicated media_type exists)
|
||||
THREE_D_EXTENSIONS = frozenset({'.obj', '.fbx', '.gltf', '.glb', '.usdz'})
|
||||
|
||||
# Text file extensions for preview fallback (the formats SaveText can produce)
|
||||
TEXT_EXTENSIONS = frozenset({'.txt', '.md', '.json'})
|
||||
|
||||
|
||||
def has_3d_extension(filename: str) -> bool:
|
||||
lower = filename.lower()
|
||||
@@ -146,10 +143,9 @@ def is_previewable(media_type: str, item: dict) -> bool:
|
||||
Maintains backwards compatibility with existing logic.
|
||||
|
||||
Priority:
|
||||
1. media_type is 'images', 'video', 'audio', '3d', or 'text'
|
||||
1. media_type is 'images', 'video', 'audio', or '3d'
|
||||
2. format field starts with 'video/' or 'audio/'
|
||||
3. filename has a 3D extension (.obj, .fbx, .gltf, .glb, .usdz)
|
||||
4. filename has a text extension (.txt, .md, .json, ...)
|
||||
"""
|
||||
if media_type in PREVIEWABLE_MEDIA_TYPES:
|
||||
return True
|
||||
@@ -160,12 +156,10 @@ def is_previewable(media_type: str, item: dict) -> bool:
|
||||
if fmt and (fmt.startswith('video/') or fmt.startswith('audio/')):
|
||||
return True
|
||||
|
||||
# Check for 3D and text files by extension
|
||||
# Check for 3D files by extension
|
||||
filename = item.get('filename', '').lower()
|
||||
if any(filename.endswith(ext) for ext in THREE_D_EXTENSIONS):
|
||||
return True
|
||||
if any(filename.endswith(ext) for ext in TEXT_EXTENSIONS):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
@@ -204,14 +198,10 @@ def normalize_history_item(prompt_id: str, history_item: dict, include_outputs:
|
||||
outputs_count, preview_output = get_outputs_summary(outputs)
|
||||
|
||||
execution_error = None
|
||||
execution_errors = []
|
||||
execution_start_time = None
|
||||
execution_end_time = None
|
||||
execution_success = None
|
||||
was_interrupted = False
|
||||
execution_summary = {}
|
||||
if status_info:
|
||||
execution_summary = status_info.get('execution_summary') or {}
|
||||
messages = status_info.get('messages', [])
|
||||
for entry in messages:
|
||||
if isinstance(entry, (list, tuple)) and len(entry) >= 2:
|
||||
@@ -221,22 +211,10 @@ def normalize_history_item(prompt_id: str, history_item: dict, include_outputs:
|
||||
execution_start_time = event_data.get('timestamp')
|
||||
elif event_name in ('execution_success', 'execution_error', 'execution_interrupted'):
|
||||
execution_end_time = event_data.get('timestamp')
|
||||
if event_name == 'execution_success':
|
||||
execution_success = event_data
|
||||
elif event_name == 'execution_error':
|
||||
if event_name == 'execution_error':
|
||||
execution_error = event_data
|
||||
elif event_name == 'execution_interrupted':
|
||||
was_interrupted = True
|
||||
elif event_name == 'execution_node_error':
|
||||
execution_errors.append(event_data)
|
||||
|
||||
completion_status = execution_summary.get('completion_status')
|
||||
if completion_status is None and execution_success is not None:
|
||||
completion_status = execution_success.get('completion_status', 'success')
|
||||
if completion_status is None and status_str == 'success':
|
||||
completion_status = 'success'
|
||||
has_errors = execution_summary.get('has_errors', bool(execution_errors)) if completion_status is not None else None
|
||||
execution_error_count = execution_summary.get('execution_error_count', len(execution_errors)) if completion_status is not None else None
|
||||
|
||||
if status_str == 'success':
|
||||
status = JobStatus.COMPLETED
|
||||
@@ -253,9 +231,6 @@ def normalize_history_item(prompt_id: str, history_item: dict, include_outputs:
|
||||
'execution_start_time': execution_start_time,
|
||||
'execution_end_time': execution_end_time,
|
||||
'execution_error': execution_error,
|
||||
'completion_status': completion_status,
|
||||
'has_errors': has_errors,
|
||||
'execution_error_count': execution_error_count,
|
||||
'outputs_count': outputs_count,
|
||||
'preview_output': preview_output,
|
||||
'workflow_id': workflow_id,
|
||||
@@ -264,7 +239,6 @@ def normalize_history_item(prompt_id: str, history_item: dict, include_outputs:
|
||||
if include_outputs:
|
||||
job['outputs'] = normalize_outputs(outputs)
|
||||
job['execution_status'] = status_info
|
||||
job['execution_errors'] = execution_errors
|
||||
job['workflow'] = {
|
||||
'prompt': prompt,
|
||||
'extra_data': extra_data,
|
||||
@@ -281,10 +255,6 @@ def get_outputs_summary(outputs: dict) -> tuple[int, Optional[dict]]:
|
||||
Preview priority (matching frontend):
|
||||
1. type="output" with previewable media
|
||||
2. Any previewable media
|
||||
|
||||
Text content entries (strings under 'text') are preview-only metadata,
|
||||
matching the frontend's METADATA_KEYS: they can serve as the fallback
|
||||
preview but are not counted as outputs.
|
||||
"""
|
||||
count = 0
|
||||
preview_output = None
|
||||
@@ -305,6 +275,7 @@ def get_outputs_summary(outputs: dict) -> tuple[int, Optional[dict]]:
|
||||
if normalized is None:
|
||||
# Not a 3D file string — check for text preview
|
||||
if media_type == 'text':
|
||||
count += 1
|
||||
if preview_output is None:
|
||||
if isinstance(item, tuple):
|
||||
text_value = item[0] if item else ''
|
||||
|
||||
@@ -17,7 +17,6 @@ class NodeState(Enum):
|
||||
Running = "running"
|
||||
Finished = "finished"
|
||||
Error = "error"
|
||||
Blocked = "blocked"
|
||||
|
||||
|
||||
class NodeProgressState(TypedDict):
|
||||
@@ -302,25 +301,17 @@ class ProgressRegistry:
|
||||
node_id, value, max_value, entry, self.prompt_id, image
|
||||
)
|
||||
|
||||
def _finish_progress(self, node_id: str, state: NodeState) -> None:
|
||||
def finish_progress(self, node_id: str) -> None:
|
||||
"""Finish progress tracking for a node"""
|
||||
entry = self.ensure_entry(node_id)
|
||||
entry["state"] = state
|
||||
entry["state"] = NodeState.Finished
|
||||
entry["value"] = entry["max"]
|
||||
|
||||
# Notify all enabled handlers
|
||||
for handler in self.handlers.values():
|
||||
if handler.enabled:
|
||||
handler.finish_handler(node_id, entry, self.prompt_id)
|
||||
|
||||
def finish_progress(self, node_id: str) -> None:
|
||||
"""Finish progress tracking for a node"""
|
||||
self._finish_progress(node_id, NodeState.Finished)
|
||||
|
||||
def error_progress(self, node_id: str) -> None:
|
||||
self._finish_progress(node_id, NodeState.Error)
|
||||
|
||||
def block_progress(self, node_id: str) -> None:
|
||||
self._finish_progress(node_id, NodeState.Blocked)
|
||||
|
||||
def reset_handlers(self) -> None:
|
||||
"""Reset all handlers"""
|
||||
for handler in self.handlers.values():
|
||||
|
||||
@@ -298,7 +298,6 @@ class PreviewAudio(IO.ComfyNode):
|
||||
search_aliases=["play audio"],
|
||||
display_name="Preview Audio",
|
||||
category="audio",
|
||||
description="Preview the audio without saving it to the ComfyUI output directory.",
|
||||
inputs=[
|
||||
IO.Audio.Input("audio"),
|
||||
],
|
||||
|
||||
@@ -1,5 +1,3 @@
|
||||
import json
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image, ImageDraw, ImageEnhance, ImageFont
|
||||
@@ -168,111 +166,6 @@ def boxes_to_regions(boxes, width: int, height: int) -> list:
|
||||
return regions
|
||||
|
||||
|
||||
def normalize_incoming_boxes(bboxes) -> list:
|
||||
if isinstance(bboxes, dict):
|
||||
frame = [bboxes]
|
||||
elif not isinstance(bboxes, list) or not bboxes:
|
||||
frame = []
|
||||
elif isinstance(bboxes[0], dict):
|
||||
frame = bboxes
|
||||
else:
|
||||
frame = bboxes[0] if isinstance(bboxes[0], list) else []
|
||||
boxes = []
|
||||
for box in frame:
|
||||
if not isinstance(box, dict):
|
||||
continue
|
||||
norm = {
|
||||
"x": box.get("x", 0),
|
||||
"y": box.get("y", 0),
|
||||
"width": box.get("width", 0),
|
||||
"height": box.get("height", 0),
|
||||
}
|
||||
meta = box.get("metadata")
|
||||
if isinstance(meta, dict):
|
||||
norm["metadata"] = meta
|
||||
boxes.append(norm)
|
||||
return boxes
|
||||
|
||||
|
||||
def _looks_like_element(box: dict) -> bool:
|
||||
bbox = box.get("bbox")
|
||||
return isinstance(bbox, (list, tuple)) and len(bbox) == 4
|
||||
|
||||
|
||||
def _looks_like_bbox(box: dict) -> bool:
|
||||
return all(key in box for key in ("x", "y", "width", "height"))
|
||||
|
||||
|
||||
def elements_to_boxes(elements: list, width: int, height: int) -> list:
|
||||
boxes = []
|
||||
for element in elements:
|
||||
if not isinstance(element, dict):
|
||||
continue
|
||||
bbox = element.get("bbox")
|
||||
if not (isinstance(bbox, (list, tuple)) and len(bbox) == 4):
|
||||
raise ValueError("bboxes element is missing a valid 'bbox' [ymin, xmin, ymax, xmax]")
|
||||
try:
|
||||
ymin, xmin, ymax, xmax = (float(v) / 1000.0 for v in bbox)
|
||||
except (TypeError, ValueError):
|
||||
raise ValueError("bboxes element 'bbox' must contain four numbers")
|
||||
etype = "text" if element.get("type") == "text" else "obj"
|
||||
boxes.append({
|
||||
"x": round(min(xmin, xmax) * width),
|
||||
"y": round(min(ymin, ymax) * height),
|
||||
"width": round(abs(xmax - xmin) * width),
|
||||
"height": round(abs(ymax - ymin) * height),
|
||||
"metadata": {
|
||||
"type": etype,
|
||||
"text": element.get("text", "") if etype == "text" else "",
|
||||
"desc": element.get("desc", ""),
|
||||
"palette": element.get("color_palette", []) or [],
|
||||
},
|
||||
})
|
||||
return boxes
|
||||
|
||||
|
||||
def boxes_from_input(data, width: int, height: int) -> list:
|
||||
if data is None:
|
||||
return []
|
||||
if isinstance(data, str):
|
||||
text = data.strip()
|
||||
if not text:
|
||||
return []
|
||||
try:
|
||||
data = json.loads(text)
|
||||
except (ValueError, TypeError) as exc:
|
||||
raise ValueError(f"bboxes string input is not valid JSON: {exc}") from exc
|
||||
if isinstance(data, dict):
|
||||
if _looks_like_element(data):
|
||||
return elements_to_boxes([data], width, height)
|
||||
if _looks_like_bbox(data):
|
||||
return normalize_incoming_boxes(data)
|
||||
raise ValueError(
|
||||
"bboxes dict must be a bounding box (x, y, width, height) or an element (with a 'bbox')"
|
||||
)
|
||||
if not isinstance(data, list):
|
||||
raise ValueError(
|
||||
"bboxes input must be bounding boxes, elements, or a JSON string, "
|
||||
f"got {type(data).__name__}"
|
||||
)
|
||||
if not data:
|
||||
return []
|
||||
first = data[0]
|
||||
if isinstance(first, list):
|
||||
return normalize_incoming_boxes(data)
|
||||
if isinstance(first, dict):
|
||||
if _looks_like_element(first):
|
||||
return elements_to_boxes(data, width, height)
|
||||
if _looks_like_bbox(first):
|
||||
return normalize_incoming_boxes(data)
|
||||
raise ValueError(
|
||||
"bboxes items must be bounding boxes (x, y, width, height) or elements (with a 'bbox')"
|
||||
)
|
||||
raise ValueError(
|
||||
f"bboxes list must contain bounding boxes or elements, got {type(first).__name__}"
|
||||
)
|
||||
|
||||
|
||||
def _norm_bbox(region: dict) -> list[int]:
|
||||
def grid(value: float) -> int:
|
||||
return max(0, min(1000, round(value * 1000)))
|
||||
@@ -324,48 +217,29 @@ class CreateBoundingBoxes(io.ComfyNode):
|
||||
optional=True,
|
||||
tooltip="Optional image used as background in the canvas and preview.",
|
||||
),
|
||||
io.MultiType.Input(
|
||||
"bboxes",
|
||||
[io.BoundingBox, io.Array, io.String],
|
||||
optional=True,
|
||||
tooltip="Bounding boxes, elements, or a JSON string to initialize the canvas. A new upstream value initializes the canvas; edits made on the canvas take priority and are kept until the upstream value changes again.",
|
||||
),
|
||||
io.Int.Input("width", default=1024, min=64, max=16384, step=16,
|
||||
tooltip="Width of the canvas and the pixel grid for the bounding boxes."),
|
||||
io.Int.Input("height", default=1024, min=64, max=16384, step=16,
|
||||
tooltip="Height of the canvas and the pixel grid for the bounding boxes."),
|
||||
editor_state,
|
||||
io.BoundingBoxes.Input(
|
||||
"last_incoming",
|
||||
optional=True,
|
||||
tooltip="Internal state managed by the canvas: the upstream bboxes value that last initialized it. Leave empty to re-initialize the canvas from the bboxes input on the next run.",
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
io.Image.Output(display_name="preview"),
|
||||
io.BoundingBox.Output(display_name="bboxes"),
|
||||
io.Array.Output(display_name="elements"),
|
||||
],
|
||||
is_output_node=True,
|
||||
is_experimental=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, width, height, editor_state=None, last_incoming=None, background=None, bboxes=None) -> io.NodeOutput:
|
||||
incoming = boxes_from_input(bboxes, width, height)
|
||||
applied = last_incoming if isinstance(last_incoming, list) else []
|
||||
upstream_changed = bool(incoming) and incoming != applied
|
||||
source = incoming if upstream_changed else (editor_state or [])
|
||||
regions = boxes_to_regions(source, width, height)
|
||||
def execute(cls, width, height, editor_state=None, background=None) -> io.NodeOutput:
|
||||
regions = boxes_to_regions(editor_state, width, height)
|
||||
preview = render_preview(regions, width, height, _bg_from_image(background))
|
||||
ui = {"dims": [width, height]}
|
||||
if incoming:
|
||||
ui["input_bboxes"] = incoming
|
||||
return io.NodeOutput(
|
||||
preview,
|
||||
fractions_to_bbox_frame(regions, width, height),
|
||||
build_elements(regions),
|
||||
ui=ui,
|
||||
ui={"dims": [width, height]},
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -844,18 +844,15 @@ class ImageMergeTileList(IO.ComfyNode):
|
||||
# Format specifications
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Maps (file_format, bit_depth, num_channels) -> (quantization scale, numpy dtype,
|
||||
# av frame pix_fmt, stream pix_fmt). Keeps the encode path declarative instead of branchy.
|
||||
# Maps (file_format, bit_depth, has_alpha) -> (numpy dtype scale, av pixel format,
|
||||
# stream pix_fmt). Keeps the encode path declarative instead of branchy.
|
||||
_FORMAT_SPECS = {
|
||||
("png", "8-bit", 1): {"scale": 255.0, "dtype": np.uint8, "frame_fmt": "gray", "stream_fmt": "gray"},
|
||||
("png", "8-bit", 3): {"scale": 255.0, "dtype": np.uint8, "frame_fmt": "rgb24", "stream_fmt": "rgb24"},
|
||||
("png", "8-bit", 4): {"scale": 255.0, "dtype": np.uint8, "frame_fmt": "rgba", "stream_fmt": "rgba"},
|
||||
("png", "16-bit", 1): {"scale": 65535.0, "dtype": np.uint16, "frame_fmt": "gray16le", "stream_fmt": "gray16be"},
|
||||
("png", "16-bit", 3): {"scale": 65535.0, "dtype": np.uint16, "frame_fmt": "rgb48le", "stream_fmt": "rgb48be"},
|
||||
("png", "16-bit", 4): {"scale": 65535.0, "dtype": np.uint16, "frame_fmt": "rgba64le", "stream_fmt": "rgba64be"},
|
||||
("exr", "32-bit float", 1): {"scale": 1.0, "dtype": np.float32, "frame_fmt": "grayf32le", "stream_fmt": "grayf32le"},
|
||||
("exr", "32-bit float", 3): {"scale": 1.0, "dtype": np.float32, "frame_fmt": "gbrpf32le", "stream_fmt": "gbrpf32le"},
|
||||
("exr", "32-bit float", 4): {"scale": 1.0, "dtype": np.float32, "frame_fmt": "gbrapf32le", "stream_fmt": "gbrapf32le"},
|
||||
("png", "8-bit", False): {"scale": 255.0, "dtype": np.uint8, "frame_fmt": "rgb24", "stream_fmt": "rgb24"},
|
||||
("png", "8-bit", True): {"scale": 255.0, "dtype": np.uint8, "frame_fmt": "rgba", "stream_fmt": "rgba"},
|
||||
("png", "16-bit", False): {"scale": 65535.0, "dtype": np.uint16, "frame_fmt": "rgb48le", "stream_fmt": "rgb48be"},
|
||||
("png", "16-bit", True): {"scale": 65535.0, "dtype": np.uint16, "frame_fmt": "rgba64le", "stream_fmt": "rgba64be"},
|
||||
("exr", "32-bit float", False): {"scale": 1.0, "dtype": np.float32, "frame_fmt": "gbrpf32le", "stream_fmt": "gbrpf32le"},
|
||||
("exr", "32-bit float", True): {"scale": 1.0, "dtype": np.float32, "frame_fmt": "gbrapf32le", "stream_fmt": "gbrapf32le"},
|
||||
}
|
||||
|
||||
|
||||
@@ -894,11 +891,10 @@ def hlg_to_linear(t: torch.Tensor) -> torch.Tensor:
|
||||
return torch.cat([hlg_to_linear(rgb), alpha], dim=-1)
|
||||
|
||||
# Piecewise: sqrt branch below 0.5, log branch above.
|
||||
# Clamp the log branch at the 0.5 branch point (not above it) so the
|
||||
# unselected lane stays finite in exp() without altering selected values;
|
||||
# Clamp inside the log branch so negative / out-of-range values don't blow up;
|
||||
# values above 1.0 are allowed and extrapolate naturally.
|
||||
low = (t ** 2) / 3.0
|
||||
high = (torch.exp((t.clamp(min=0.5) - _HLG_C) / _HLG_A) + _HLG_B) / 12.0
|
||||
high = (torch.exp((t.clamp(min=_HLG_C) - _HLG_C) / _HLG_A) + _HLG_B) / 12.0
|
||||
return torch.where(t <= 0.5, low, high)
|
||||
|
||||
|
||||
@@ -1091,8 +1087,7 @@ def _encode_image(
|
||||
bit_depth: str,
|
||||
colorspace: str,
|
||||
) -> bytes:
|
||||
"""Encode a single HxWxC (or channel-less HxW grayscale) tensor to PNG or
|
||||
EXR bytes in memory. Grayscale is written as single-channel PNG / Y-only EXR.
|
||||
"""Encode a single HxWxC tensor to PNG or EXR bytes in memory.
|
||||
|
||||
For EXR the input is interpreted according to `colorspace` and converted
|
||||
to scene-linear (EXR's convention) before writing:
|
||||
@@ -1106,16 +1101,10 @@ def _encode_image(
|
||||
For PNG, colorspace selection does not modify pixels — PNG is delivered
|
||||
sRGB-encoded and there is no PNG path for wide-gamut HDR in this node.
|
||||
"""
|
||||
if img_tensor.ndim == 2:
|
||||
img_tensor = img_tensor.unsqueeze(-1) # Some nodes emit grayscale as (H, W) with no channel dim, mask-style.
|
||||
height, width, num_channels = img_tensor.shape
|
||||
has_alpha = num_channels == 4
|
||||
|
||||
spec = _FORMAT_SPECS.get((file_format, bit_depth, num_channels))
|
||||
if spec is None:
|
||||
raise ValueError(
|
||||
f"No {file_format}/{bit_depth} encoder for {num_channels}-channel images: "
|
||||
"supported channel counts are 1 (grayscale), 3 (RGB) and 4 (RGBA)."
|
||||
)
|
||||
spec = _FORMAT_SPECS[(file_format, bit_depth, has_alpha)]
|
||||
|
||||
if spec["dtype"] == np.float32:
|
||||
# EXR path: preserve full range, no clamp.
|
||||
|
||||
@@ -61,10 +61,14 @@ class Load3D(IO.ComfyNode):
|
||||
|
||||
@classmethod
|
||||
def execute(cls, model_file, image, **kwargs) -> IO.NodeOutput:
|
||||
image_path = folder_paths.get_annotated_filepath(image['image'])
|
||||
mask_path = folder_paths.get_annotated_filepath(image['mask'])
|
||||
normal_path = folder_paths.get_annotated_filepath(image['normal'])
|
||||
|
||||
load_image_node = nodes.LoadImage()
|
||||
output_image, ignore_mask = load_image_node.load_image(image=image['image'])
|
||||
ignore_image, output_mask = load_image_node.load_image(image=image['mask'])
|
||||
normal_image, ignore_mask2 = load_image_node.load_image(image=image['normal'])
|
||||
output_image, ignore_mask = load_image_node.load_image(image=image_path)
|
||||
ignore_image, output_mask = load_image_node.load_image(image=mask_path)
|
||||
normal_image, ignore_mask2 = load_image_node.load_image(image=normal_path)
|
||||
|
||||
video = None
|
||||
|
||||
@@ -92,7 +96,6 @@ class Preview3D(IO.ComfyNode):
|
||||
search_aliases=["view mesh", "3d viewer"],
|
||||
display_name="Preview 3D & Animation",
|
||||
category="3d",
|
||||
description="Preview a 3D model file without saving it to the ComfyUI output directory.",
|
||||
is_experimental=True,
|
||||
is_output_node=True,
|
||||
inputs=[
|
||||
@@ -137,7 +140,6 @@ class Preview3DAdvanced(IO.ComfyNode):
|
||||
display_name="Preview 3D (Advanced)",
|
||||
search_aliases=["preview 3d", "3d viewer", "view mesh", "frame 3d", "3d camera output"],
|
||||
category="3d",
|
||||
description="Preview a 3D model file without saving it to the ComfyUI output directory.",
|
||||
is_experimental=True,
|
||||
is_output_node=True,
|
||||
inputs=[
|
||||
@@ -174,9 +176,8 @@ class Preview3DAdvanced(IO.ComfyNode):
|
||||
filename = f"preview3d_advanced_{uuid.uuid4().hex}.{model_3d.format}"
|
||||
model_3d.save_to(os.path.join(folder_paths.get_temp_directory(), filename))
|
||||
|
||||
viewport_state = viewport_state if isinstance(viewport_state, dict) else {}
|
||||
camera_info_input = kwargs.get("camera_info", None)
|
||||
camera_info = camera_info_input if camera_info_input is not None else viewport_state.get('camera_info')
|
||||
camera_info = camera_info_input if camera_info_input is not None else viewport_state['camera_info']
|
||||
model_3d_info_input = kwargs.get("model_3d_info", None)
|
||||
model_3d_info = model_3d_info_input if model_3d_info_input is not None else viewport_state.get('model_3d_info', [])
|
||||
return IO.NodeOutput(
|
||||
@@ -196,7 +197,6 @@ class PreviewGaussianSplat(IO.ComfyNode):
|
||||
node_id="PreviewGaussianSplat",
|
||||
display_name="Preview Splat",
|
||||
category="3d",
|
||||
description="Preview a gaussian splat 3D file without saving it to the ComfyUI output directory.",
|
||||
is_experimental=True,
|
||||
is_output_node=True,
|
||||
search_aliases=[
|
||||
@@ -244,9 +244,8 @@ class PreviewGaussianSplat(IO.ComfyNode):
|
||||
filename = f"preview_splat_{uuid.uuid4().hex}.{model_3d.format}"
|
||||
model_3d.save_to(os.path.join(folder_paths.get_temp_directory(), filename))
|
||||
|
||||
viewport_state = viewport_state if isinstance(viewport_state, dict) else {}
|
||||
camera_info_input = kwargs.get("camera_info", None)
|
||||
camera_info = camera_info_input if camera_info_input is not None else viewport_state.get('camera_info')
|
||||
camera_info = camera_info_input if camera_info_input is not None else viewport_state['camera_info']
|
||||
model_3d_info_input = kwargs.get("model_3d_info", None)
|
||||
model_3d_info = model_3d_info_input if model_3d_info_input is not None else viewport_state.get('model_3d_info', [])
|
||||
return IO.NodeOutput(
|
||||
@@ -266,7 +265,6 @@ class PreviewPointCloud(IO.ComfyNode):
|
||||
node_id="PreviewPointCloud",
|
||||
display_name="Preview Point Cloud",
|
||||
category="3d",
|
||||
description="Preview a point cloud 3D file without saving it to the ComfyUI output directory.",
|
||||
is_experimental=True,
|
||||
is_output_node=True,
|
||||
search_aliases=[
|
||||
@@ -305,9 +303,8 @@ class PreviewPointCloud(IO.ComfyNode):
|
||||
filename = f"preview_pointcloud_{uuid.uuid4().hex}.{model_3d.format}"
|
||||
model_3d.save_to(os.path.join(folder_paths.get_temp_directory(), filename))
|
||||
|
||||
viewport_state = viewport_state if isinstance(viewport_state, dict) else {}
|
||||
camera_info_input = kwargs.get("camera_info", None)
|
||||
camera_info = camera_info_input if camera_info_input is not None else viewport_state.get('camera_info')
|
||||
camera_info = camera_info_input if camera_info_input is not None else viewport_state['camera_info']
|
||||
model_3d_info_input = kwargs.get("model_3d_info", None)
|
||||
model_3d_info = model_3d_info_input if model_3d_info_input is not None else viewport_state.get('model_3d_info', [])
|
||||
return IO.NodeOutput(
|
||||
@@ -378,9 +375,8 @@ class Load3DAdvanced(IO.ComfyNode):
|
||||
file_3d = None
|
||||
if model_file and model_file != "none":
|
||||
file_3d = Types.File3D(folder_paths.get_annotated_filepath(model_file))
|
||||
viewport_state = viewport_state if isinstance(viewport_state, dict) else {}
|
||||
model_3d_info = viewport_state.get('model_3d_info', [])
|
||||
return IO.NodeOutput(file_3d, model_3d_info, viewport_state.get('camera_info'), width, height)
|
||||
return IO.NodeOutput(file_3d, model_3d_info, viewport_state['camera_info'], width, height)
|
||||
|
||||
|
||||
class Load3DExtension(ComfyExtension):
|
||||
|
||||
@@ -419,18 +419,17 @@ class MaskPreview(IO.ComfyNode):
|
||||
search_aliases=["show mask", "view mask", "inspect mask", "debug mask"],
|
||||
display_name="Preview Mask",
|
||||
category="image/mask",
|
||||
description="Preview the masks without saving them to the ComfyUI output directory.",
|
||||
description="Saves the input images to your ComfyUI output directory.",
|
||||
inputs=[
|
||||
IO.Mask.Input("mask"),
|
||||
],
|
||||
hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo],
|
||||
is_output_node=True,
|
||||
outputs=[IO.Mask.Output(display_name="mask")]
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, mask, filename_prefix="ComfyUI") -> IO.NodeOutput:
|
||||
return IO.NodeOutput(mask, ui=UI.PreviewMask(mask))
|
||||
return IO.NodeOutput(ui=UI.PreviewMask(mask))
|
||||
|
||||
|
||||
class MaskExtension(ComfyExtension):
|
||||
|
||||
@@ -18,7 +18,6 @@ class PreviewAny():
|
||||
|
||||
CATEGORY = "utilities"
|
||||
SEARCH_ALIASES = ["show output", "inspect", "debug", "print value", "show text"]
|
||||
DESCRIPTION = "Preview any input value as text."
|
||||
|
||||
def main(self, source=None):
|
||||
torch.set_printoptions(edgeitems=6)
|
||||
|
||||
@@ -10,10 +10,11 @@ class String(io.ComfyNode):
|
||||
return io.Schema(
|
||||
node_id="PrimitiveString",
|
||||
search_aliases=["text", "string", "text box", "prompt"],
|
||||
display_name="Text",
|
||||
display_name="Text String (DEPRECATED)",
|
||||
category="utilities/primitive",
|
||||
inputs=[io.String.Input("value")],
|
||||
outputs=[io.String.Output()]
|
||||
outputs=[io.String.Output()],
|
||||
is_deprecated=True
|
||||
)
|
||||
|
||||
@classmethod
|
||||
@@ -27,7 +28,7 @@ class StringMultiline(io.ComfyNode):
|
||||
return io.Schema(
|
||||
node_id="PrimitiveStringMultiline",
|
||||
search_aliases=["text", "string", "text multiline", "string multiline", "text box", "prompt"],
|
||||
display_name="Text (Multiline)",
|
||||
display_name="Input Text",
|
||||
category="utilities/primitive",
|
||||
essentials_category="Basics",
|
||||
inputs=[io.String.Input("value", multiline=True)],
|
||||
|
||||
@@ -13,7 +13,7 @@ from typing_extensions import override
|
||||
|
||||
import folder_paths
|
||||
from comfy.cli_args import args
|
||||
from comfy_api.latest import ComfyExtension, IO, Types, UI
|
||||
from comfy_api.latest import ComfyExtension, IO, Types
|
||||
|
||||
|
||||
def pack_variable_mesh_batch(vertices, faces, colors=None, uvs=None, texture=None, unlit=False):
|
||||
@@ -406,165 +406,10 @@ class SaveGLB(IO.ComfyNode):
|
||||
return IO.NodeOutput(ui={"3d": results})
|
||||
|
||||
|
||||
def _save_file3d_to_output(model_3d: Types.File3D, filename_prefix: str) -> str:
|
||||
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(
|
||||
filename_prefix, folder_paths.get_output_directory()
|
||||
)
|
||||
ext = model_3d.format or "glb"
|
||||
saved_filename = f"{filename}_{counter:05}.{ext}"
|
||||
model_3d.save_to(os.path.join(full_output_folder, saved_filename))
|
||||
return f"{subfolder}/{saved_filename}" if subfolder else saved_filename
|
||||
|
||||
|
||||
def execute_save_3d_advanced(model_3d, viewport_state, width, height, filename_prefix, kwargs) -> IO.NodeOutput:
|
||||
model_file = _save_file3d_to_output(model_3d, filename_prefix)
|
||||
viewport_state = viewport_state if isinstance(viewport_state, dict) else {}
|
||||
camera_info_input = kwargs.get("camera_info", None)
|
||||
camera_info = camera_info_input if camera_info_input is not None else viewport_state.get('camera_info')
|
||||
model_3d_info_input = kwargs.get("model_3d_info", None)
|
||||
model_3d_info = model_3d_info_input if model_3d_info_input is not None else viewport_state.get('model_3d_info', [])
|
||||
return IO.NodeOutput(
|
||||
model_3d,
|
||||
model_3d_info,
|
||||
camera_info,
|
||||
width,
|
||||
height,
|
||||
ui=UI.PreviewUI3DAdvanced(model_file, camera_info, model_3d_info),
|
||||
)
|
||||
|
||||
|
||||
class Save3DAdvanced(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="Save3DAdvanced",
|
||||
display_name="Save 3D (Advanced)",
|
||||
search_aliases=["save 3d", "export 3d model", "save mesh advanced"],
|
||||
category="3d",
|
||||
is_experimental=True,
|
||||
is_output_node=True,
|
||||
inputs=[
|
||||
IO.MultiType.Input(
|
||||
"model_3d",
|
||||
types=[
|
||||
IO.File3DGLB,
|
||||
IO.File3DGLTF,
|
||||
IO.File3DFBX,
|
||||
IO.File3DOBJ,
|
||||
IO.File3DSTL,
|
||||
IO.File3DUSDZ,
|
||||
IO.File3DAny,
|
||||
],
|
||||
tooltip="3D model file from an upstream 3D node.",
|
||||
),
|
||||
IO.String.Input("filename_prefix", default="3d/ComfyUI"),
|
||||
IO.Load3D.Input("viewport_state"),
|
||||
IO.Load3DModelInfo.Input("model_3d_info", optional=True, advanced=True),
|
||||
IO.Load3DCamera.Input("camera_info", optional=True, advanced=True),
|
||||
IO.Int.Input("width", default=1024, min=1, max=4096, step=1),
|
||||
IO.Int.Input("height", default=1024, min=1, max=4096, step=1),
|
||||
],
|
||||
outputs=[
|
||||
IO.File3DAny.Output(display_name="model_3d"),
|
||||
IO.Load3DModelInfo.Output(display_name="model_3d_info"),
|
||||
IO.Load3DCamera.Output(display_name="camera_info"),
|
||||
IO.Int.Output(display_name="width"),
|
||||
IO.Int.Output(display_name="height"),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, model_3d: Types.File3D, viewport_state, width: int, height: int, filename_prefix: str, **kwargs) -> IO.NodeOutput:
|
||||
return execute_save_3d_advanced(model_3d, viewport_state, width, height, filename_prefix, kwargs)
|
||||
|
||||
|
||||
class SaveGaussianSplat(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="SaveGaussianSplat",
|
||||
display_name="Save Splat",
|
||||
search_aliases=["save splat", "save gaussian splat", "export gaussian", "export splat"],
|
||||
category="3d",
|
||||
is_experimental=True,
|
||||
is_output_node=True,
|
||||
inputs=[
|
||||
IO.MultiType.Input(
|
||||
"model_3d",
|
||||
types=[
|
||||
IO.File3DSplatAny,
|
||||
IO.File3DPLY,
|
||||
IO.File3DSPLAT,
|
||||
IO.File3DSPZ,
|
||||
IO.File3DKSPLAT,
|
||||
],
|
||||
tooltip="A gaussian splat 3D file.",
|
||||
),
|
||||
IO.String.Input("filename_prefix", default="3d/ComfyUI"),
|
||||
IO.Load3D.Input("viewport_state"),
|
||||
IO.Load3DModelInfo.Input("model_3d_info", optional=True, advanced=True),
|
||||
IO.Load3DCamera.Input("camera_info", optional=True, advanced=True),
|
||||
IO.Int.Input("width", default=1024, min=1, max=4096, step=1),
|
||||
IO.Int.Input("height", default=1024, min=1, max=4096, step=1),
|
||||
],
|
||||
outputs=[
|
||||
IO.File3DSplatAny.Output(display_name="model_3d"),
|
||||
IO.Load3DModelInfo.Output(display_name="model_3d_info"),
|
||||
IO.Load3DCamera.Output(display_name="camera_info"),
|
||||
IO.Int.Output(display_name="width"),
|
||||
IO.Int.Output(display_name="height"),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, model_3d: Types.File3D, viewport_state, width: int, height: int, filename_prefix: str, **kwargs) -> IO.NodeOutput:
|
||||
return execute_save_3d_advanced(model_3d, viewport_state, width, height, filename_prefix, kwargs)
|
||||
|
||||
|
||||
class SavePointCloud(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="SavePointCloud",
|
||||
display_name="Save Point Cloud",
|
||||
search_aliases=["save point cloud", "save pointcloud", "export point cloud"],
|
||||
category="3d",
|
||||
is_experimental=True,
|
||||
is_output_node=True,
|
||||
inputs=[
|
||||
IO.MultiType.Input(
|
||||
"model_3d",
|
||||
types=[
|
||||
IO.File3DPointCloudAny,
|
||||
IO.File3DPLY,
|
||||
],
|
||||
tooltip="Point cloud file (.ply)",
|
||||
),
|
||||
IO.String.Input("filename_prefix", default="3d/ComfyUI"),
|
||||
IO.Load3D.Input("viewport_state"),
|
||||
IO.Load3DModelInfo.Input("model_3d_info", optional=True, advanced=True),
|
||||
IO.Load3DCamera.Input("camera_info", optional=True, advanced=True),
|
||||
IO.Int.Input("width", default=1024, min=1, max=4096, step=1),
|
||||
IO.Int.Input("height", default=1024, min=1, max=4096, step=1),
|
||||
],
|
||||
outputs=[
|
||||
IO.File3DPointCloudAny.Output(display_name="model_3d"),
|
||||
IO.Load3DModelInfo.Output(display_name="model_3d_info"),
|
||||
IO.Load3DCamera.Output(display_name="camera_info"),
|
||||
IO.Int.Output(display_name="width"),
|
||||
IO.Int.Output(display_name="height"),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, model_3d: Types.File3D, viewport_state, width: int, height: int, filename_prefix: str, **kwargs) -> IO.NodeOutput:
|
||||
return execute_save_3d_advanced(model_3d, viewport_state, width, height, filename_prefix, kwargs)
|
||||
|
||||
|
||||
class Save3DExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
|
||||
return [SaveGLB, Save3DAdvanced, SaveGaussianSplat, SavePointCloud]
|
||||
return [SaveGLB]
|
||||
|
||||
|
||||
async def comfy_entrypoint() -> Save3DExtension:
|
||||
|
||||
@@ -1,614 +0,0 @@
|
||||
import logging
|
||||
|
||||
from typing_extensions import override
|
||||
from comfy_api.latest import ComfyExtension, io
|
||||
import torch
|
||||
|
||||
import comfy.model_management
|
||||
from comfy.ldm.seedvr.color_fix import (
|
||||
adain_color_transfer,
|
||||
lab_color_transfer,
|
||||
wavelet_color_transfer,
|
||||
)
|
||||
from comfy.ldm.seedvr.constants import (
|
||||
BYTEDANCE_VAE_SPATIAL_DOWNSAMPLE,
|
||||
SEEDVR2_ADAIN_SCALE_MULTIPLIER,
|
||||
SEEDVR2_CHUNK_GIB_PER_MPX_FRAME,
|
||||
SEEDVR2_CHUNK_RESERVED_GIB,
|
||||
SEEDVR2_CHUNK_SIGMA_GIB,
|
||||
SEEDVR2_CHUNK_SIGMA_K,
|
||||
SEEDVR2_COLOR_MEM_HEADROOM,
|
||||
SEEDVR2_DTYPE_BYTES_FLOOR,
|
||||
SEEDVR2_LAB_SCALE_MULTIPLIER,
|
||||
SEEDVR2_LATENT_CHANNELS,
|
||||
SEEDVR2_OOM_BACKOFF_DIVISOR,
|
||||
SEEDVR2_WAVELET_SCALE_MULTIPLIER,
|
||||
)
|
||||
|
||||
from torchvision.transforms import functional as TVF
|
||||
from torchvision.transforms.functional import InterpolationMode
|
||||
|
||||
|
||||
_SEEDVR2_INVALID_MODEL_MSG_PREFIX = "SeedVR2Conditioning: model object does not match expected SeedVR2 structure"
|
||||
_ATTR_MISSING = object()
|
||||
|
||||
|
||||
def _resolve_seedvr2_diffusion_model(model):
|
||||
inner = getattr(model, "model", _ATTR_MISSING)
|
||||
if inner is _ATTR_MISSING:
|
||||
raise RuntimeError(
|
||||
f"{_SEEDVR2_INVALID_MODEL_MSG_PREFIX}: input has no 'model' attribute "
|
||||
f"(got type {type(model).__name__})."
|
||||
)
|
||||
if inner is None:
|
||||
raise RuntimeError(
|
||||
f"{_SEEDVR2_INVALID_MODEL_MSG_PREFIX}: input.model is None "
|
||||
f"(input type {type(model).__name__})."
|
||||
)
|
||||
diffusion_model = getattr(inner, "diffusion_model", _ATTR_MISSING)
|
||||
if diffusion_model is _ATTR_MISSING:
|
||||
raise RuntimeError(
|
||||
f"{_SEEDVR2_INVALID_MODEL_MSG_PREFIX}: 'model.model' has no "
|
||||
f"'diffusion_model' attribute (got type {type(inner).__name__})."
|
||||
)
|
||||
if diffusion_model is None:
|
||||
raise RuntimeError(
|
||||
f"{_SEEDVR2_INVALID_MODEL_MSG_PREFIX}: 'model.model.diffusion_model' "
|
||||
f"is None (model.model type {type(inner).__name__})."
|
||||
)
|
||||
return diffusion_model
|
||||
|
||||
|
||||
def div_pad(image, factor):
|
||||
height_factor, width_factor = factor
|
||||
height, width = image.shape[-2:]
|
||||
|
||||
pad_height = (height_factor - (height % height_factor)) % height_factor
|
||||
pad_width = (width_factor - (width % width_factor)) % width_factor
|
||||
|
||||
if pad_height == 0 and pad_width == 0:
|
||||
return image
|
||||
|
||||
padding = (0, pad_width, 0, pad_height)
|
||||
return torch.nn.functional.pad(image, padding, mode='constant', value=0.0)
|
||||
|
||||
def cut_videos(videos):
|
||||
t = videos.size(1)
|
||||
if t < 1:
|
||||
raise ValueError("SeedVR2Preprocess expected at least one frame.")
|
||||
if t == 1:
|
||||
return videos
|
||||
if t <= 4:
|
||||
padding = videos[:, -1:].repeat(1, 4 - t + 1, 1, 1, 1)
|
||||
return torch.cat([videos, padding], dim=1)
|
||||
if (t - 1) % 4 == 0:
|
||||
return videos
|
||||
padding = videos[:, -1:].repeat(1, 4 - ((t - 1) % 4), 1, 1, 1)
|
||||
videos = torch.cat([videos, padding], dim=1)
|
||||
if (videos.size(1) - 1) % 4 != 0:
|
||||
raise ValueError(f"SeedVR2Preprocess failed to pad video length to 4n+1; got {videos.size(1)} frames.")
|
||||
return videos
|
||||
|
||||
def _seedvr2_input_shorter_edge(images, node_name):
|
||||
if images.dim() == 4:
|
||||
return min(images.shape[1], images.shape[2])
|
||||
if images.dim() == 5:
|
||||
return min(images.shape[2], images.shape[3])
|
||||
raise ValueError(
|
||||
f"{node_name}: expected 4-D or 5-D IMAGE tensor, "
|
||||
f"got shape {tuple(images.shape)}"
|
||||
)
|
||||
|
||||
|
||||
def _seedvr2_pad(images, upscaled_shorter_edge, node_name):
|
||||
if upscaled_shorter_edge < 2:
|
||||
raise ValueError(
|
||||
f"{node_name}: input shorter edge must be at least 2 pixels; "
|
||||
f"got {upscaled_shorter_edge}."
|
||||
)
|
||||
if images.shape[-1] > 3:
|
||||
images = images[..., :3]
|
||||
if images.dim() == 4:
|
||||
# Comfy video components arrive as a 4-D IMAGE frame sequence:
|
||||
# (frames, H, W, C). SeedVR2 consumes that as one video.
|
||||
images = images.unsqueeze(0)
|
||||
elif images.dim() != 5:
|
||||
raise ValueError(
|
||||
f"{node_name}: expected 4-D or 5-D IMAGE tensor, "
|
||||
f"got shape {tuple(images.shape)}"
|
||||
)
|
||||
images = images.permute(0, 1, 4, 2, 3)
|
||||
|
||||
b, t, c, h, w = images.shape
|
||||
images = images.reshape(b * t, c, h, w)
|
||||
|
||||
images = torch.clamp(images, 0.0, 1.0)
|
||||
images = div_pad(images, (16, 16))
|
||||
_, _, new_h, new_w = images.shape
|
||||
|
||||
images = images.reshape(b, t, c, new_h, new_w)
|
||||
images = cut_videos(images)
|
||||
images_bthwc = images.permute(0, 1, 3, 4, 2).contiguous()
|
||||
|
||||
return io.NodeOutput(images_bthwc)
|
||||
|
||||
|
||||
class SeedVR2Preprocess(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="SeedVR2Preprocess",
|
||||
display_name="Pre-Process SeedVR2 Input",
|
||||
category="image/pre-processors",
|
||||
description="Pad a resized image for SeedVR2 model. Alpha channel is dropped. The node Post-Process SeedVR2 Output re-applies it from the original resized image.",
|
||||
search_aliases=["seedvr2", "upscale", "video upscale", "pad", "preprocess"],
|
||||
inputs=[
|
||||
io.Image.Input("resized_images", tooltip="The resized image to process."),
|
||||
],
|
||||
outputs=[
|
||||
io.Image.Output("images", tooltip="The padded image for VAE encoding."),
|
||||
]
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, resized_images):
|
||||
upscaled_shorter_edge = _seedvr2_input_shorter_edge(resized_images, "SeedVR2Preprocess")
|
||||
return _seedvr2_pad(
|
||||
resized_images, upscaled_shorter_edge, "SeedVR2Preprocess",
|
||||
)
|
||||
|
||||
|
||||
class SeedVR2PostProcessing(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="SeedVR2PostProcessing",
|
||||
display_name="Post-Process SeedVR2 Output",
|
||||
category="image/post-processors",
|
||||
description="Align the generated image with the original resized image and apply color correction.",
|
||||
search_aliases=["seedvr2", "upscale", "color correction", "color match", "postprocess"],
|
||||
inputs=[
|
||||
io.Image.Input("images", tooltip="The generated image to process."),
|
||||
io.Image.Input("original_resized_images", tooltip="The original resized image before pre-processing, used as reference."),
|
||||
io.Combo.Input("color_correction_method", options=["lab", "wavelet", "adain", "none"], default="lab", tooltip="Method to match the generated image colors to the original image. lab: transfer color in CIELAB space, preserving detail (most faithful). wavelet: transfer low-frequency color, keeping upscaled high-frequency detail. adain: match per-channel mean/std (fastest, global tint). none: skip color transfer (geometry alignment only)."),
|
||||
],
|
||||
outputs=[io.Image.Output(display_name="images", tooltip="The aligned, color-corrected image.")],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, images, original_resized_images, color_correction_method):
|
||||
alpha_input = None
|
||||
if original_resized_images.shape[-1] == 4:
|
||||
alpha_input = original_resized_images[..., 3:4]
|
||||
original_resized_images = original_resized_images[..., :3]
|
||||
decoded_5d, decoded_was_4d = cls._as_bthwc(images)
|
||||
reference_full, _ = cls._as_bthwc(original_resized_images)
|
||||
decoded_5d = cls._restore_reference_batch_time(decoded_5d, reference_full)
|
||||
|
||||
b = min(decoded_5d.shape[0], reference_full.shape[0])
|
||||
t = min(decoded_5d.shape[1], reference_full.shape[1])
|
||||
reference_h = reference_full.shape[2]
|
||||
reference_w = reference_full.shape[3]
|
||||
|
||||
decoded_5d = decoded_5d[:b, :t, :, :, :]
|
||||
target_h = min(decoded_5d.shape[2], reference_h)
|
||||
target_w = min(decoded_5d.shape[3], reference_w)
|
||||
decoded_5d = decoded_5d[:, :, :target_h, :target_w, :]
|
||||
if color_correction_method in ("lab", "wavelet", "adain"):
|
||||
reference_5d = reference_full[:b, :t, :, :, :]
|
||||
reference_5d = cls._resize_reference(reference_5d, target_h, target_w)
|
||||
output_device = decoded_5d.device
|
||||
decoded_raw = cls._to_seedvr2_raw(decoded_5d)
|
||||
reference_raw = cls._to_seedvr2_raw(reference_5d)
|
||||
decoded_flat = decoded_raw.permute(0, 1, 4, 2, 3).reshape(b * t, decoded_raw.shape[4], target_h, target_w)
|
||||
reference_flat = reference_raw.permute(0, 1, 4, 2, 3).reshape(b * t, reference_raw.shape[4], target_h, target_w)
|
||||
output = cls._color_transfer_chunked(
|
||||
decoded_flat, reference_flat, output_device, color_correction_method,
|
||||
)
|
||||
output = output.reshape(b, t, output.shape[1], output.shape[2], output.shape[3]).permute(0, 1, 3, 4, 2)
|
||||
output = output.add(1.0).div(2.0).clamp(0.0, 1.0)
|
||||
elif color_correction_method == "none":
|
||||
output = decoded_5d
|
||||
else:
|
||||
raise ValueError(f"SeedVR2PostProcessing: unknown color_correction_method {color_correction_method!r}")
|
||||
|
||||
if alpha_input is not None:
|
||||
alpha_5d, _ = cls._as_bthwc(alpha_input)
|
||||
alpha_5d = alpha_5d[:output.shape[0], :output.shape[1], :output.shape[2], :output.shape[3], :]
|
||||
output = torch.cat([output, alpha_5d.to(dtype=output.dtype, device=output.device)], dim=-1)
|
||||
h2 = output.shape[-3] - (output.shape[-3] % 2)
|
||||
w2 = output.shape[-2] - (output.shape[-2] % 2)
|
||||
output = output[:, :, :h2, :w2, :]
|
||||
if decoded_was_4d:
|
||||
output = output.reshape(-1, output.shape[-3], output.shape[-2], output.shape[-1])
|
||||
return io.NodeOutput(output)
|
||||
|
||||
@staticmethod
|
||||
def _as_bthwc(images):
|
||||
if images.ndim == 4:
|
||||
return images.unsqueeze(0), True
|
||||
if images.ndim == 5:
|
||||
return images, False
|
||||
raise ValueError(
|
||||
f"SeedVR2PostProcessing: expected 4-D or 5-D IMAGE tensor, got shape {tuple(images.shape)}"
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _restore_reference_batch_time(decoded, reference):
|
||||
if decoded.shape[0] != 1:
|
||||
return decoded
|
||||
ref_b, ref_t = reference.shape[:2]
|
||||
if ref_b < 1 or decoded.shape[1] % ref_b != 0:
|
||||
return decoded
|
||||
decoded_t = decoded.shape[1] // ref_b
|
||||
if decoded_t < ref_t:
|
||||
return decoded
|
||||
return decoded.reshape(ref_b, decoded_t, decoded.shape[2], decoded.shape[3], decoded.shape[4])
|
||||
|
||||
@staticmethod
|
||||
def _to_seedvr2_raw(images):
|
||||
return images.mul(2.0).sub(1.0)
|
||||
|
||||
@staticmethod
|
||||
def _color_transfer_on_vae_device(decoded_flat, reference_flat, output_device, transfer_fn):
|
||||
color_device = comfy.model_management.vae_device()
|
||||
decoded_flat = decoded_flat.to(device=color_device)
|
||||
reference_flat = reference_flat.to(device=color_device)
|
||||
output = transfer_fn(decoded_flat, reference_flat)
|
||||
return output.to(device=output_device)
|
||||
|
||||
@staticmethod
|
||||
def _lab_color_transfer_on_vae_device(decoded_flat, reference_flat, output_device):
|
||||
color_device = comfy.model_management.vae_device()
|
||||
result = None
|
||||
for start in range(decoded_flat.shape[0]):
|
||||
decoded_frame = decoded_flat[start:start + 1].to(device=color_device).clone()
|
||||
reference_frame = reference_flat[start:start + 1].to(device=color_device).clone()
|
||||
output = lab_color_transfer(decoded_frame, reference_frame).to(device=output_device)
|
||||
if result is None:
|
||||
result = torch.empty(
|
||||
(decoded_flat.shape[0],) + tuple(output.shape[1:]),
|
||||
device=output_device,
|
||||
dtype=output.dtype,
|
||||
)
|
||||
result[start:start + 1].copy_(output)
|
||||
if result is None:
|
||||
raise ValueError("SeedVR2PostProcessing: LAB color correction requires at least one frame.")
|
||||
return result
|
||||
|
||||
@classmethod
|
||||
def _color_transfer_chunked(cls, decoded_flat, reference_flat, output_device, color_correction_method):
|
||||
chunk_size = cls._estimate_color_correction_chunk_size(decoded_flat, color_correction_method)
|
||||
while True:
|
||||
try:
|
||||
return cls._run_color_transfer_chunks(
|
||||
decoded_flat, reference_flat, output_device, color_correction_method, chunk_size,
|
||||
)
|
||||
except Exception as e:
|
||||
comfy.model_management.raise_non_oom(e)
|
||||
if chunk_size <= 1:
|
||||
raise RuntimeError(
|
||||
"SeedVR2PostProcessing: color correction OOM at one frame; "
|
||||
f"color_correction_method={color_correction_method}, shape={tuple(decoded_flat.shape)}."
|
||||
) from e
|
||||
chunk_size = max(1, chunk_size // SEEDVR2_OOM_BACKOFF_DIVISOR)
|
||||
|
||||
@classmethod
|
||||
def _run_color_transfer_chunks(cls, decoded_flat, reference_flat, output_device, color_correction_method, chunk_size):
|
||||
result = None
|
||||
for start in range(0, decoded_flat.shape[0], chunk_size):
|
||||
end = min(start + chunk_size, decoded_flat.shape[0])
|
||||
decoded_chunk = decoded_flat[start:end]
|
||||
reference_chunk = reference_flat[start:end]
|
||||
if color_correction_method == "lab":
|
||||
output = cls._lab_color_transfer_on_vae_device(decoded_chunk, reference_chunk, output_device)
|
||||
elif color_correction_method == "wavelet":
|
||||
output = cls._color_transfer_on_vae_device(
|
||||
decoded_chunk, reference_chunk, output_device, wavelet_color_transfer,
|
||||
)
|
||||
else:
|
||||
output = cls._color_transfer_on_vae_device(
|
||||
decoded_chunk, reference_chunk, output_device, adain_color_transfer,
|
||||
)
|
||||
if result is None:
|
||||
result = torch.empty(
|
||||
(decoded_flat.shape[0],) + tuple(output.shape[1:]),
|
||||
device=output_device,
|
||||
dtype=output.dtype,
|
||||
)
|
||||
result[start:end].copy_(output)
|
||||
if result is None:
|
||||
raise ValueError("SeedVR2PostProcessing: color correction requires at least one frame.")
|
||||
return result
|
||||
|
||||
@classmethod
|
||||
def _estimate_color_correction_chunk_size(cls, decoded_flat, color_correction_method):
|
||||
multiplier = cls._color_correction_memory_multiplier(color_correction_method)
|
||||
frames = decoded_flat.shape[0]
|
||||
_, channels, height, width = decoded_flat.shape
|
||||
dtype_bytes = max(decoded_flat.element_size(), SEEDVR2_DTYPE_BYTES_FLOOR)
|
||||
bytes_per_frame = height * width * channels * dtype_bytes * multiplier
|
||||
if bytes_per_frame <= 0:
|
||||
return frames
|
||||
color_device = comfy.model_management.vae_device()
|
||||
free_memory = comfy.model_management.get_free_memory(color_device)
|
||||
chunk_size = int((free_memory * SEEDVR2_COLOR_MEM_HEADROOM) // bytes_per_frame)
|
||||
return max(1, min(frames, chunk_size))
|
||||
|
||||
@staticmethod
|
||||
def _color_correction_memory_multiplier(color_correction_method):
|
||||
if color_correction_method == "lab":
|
||||
return SEEDVR2_LAB_SCALE_MULTIPLIER
|
||||
if color_correction_method == "wavelet":
|
||||
return SEEDVR2_WAVELET_SCALE_MULTIPLIER
|
||||
if color_correction_method == "adain":
|
||||
return SEEDVR2_ADAIN_SCALE_MULTIPLIER
|
||||
raise ValueError(f"SeedVR2PostProcessing: unknown color_correction_method {color_correction_method!r}")
|
||||
|
||||
@staticmethod
|
||||
def _resize_reference(reference, height, width):
|
||||
if reference.shape[2] == height and reference.shape[3] == width:
|
||||
return reference
|
||||
b, t = reference.shape[:2]
|
||||
reference_flat = reference.permute(0, 1, 4, 2, 3).reshape(b * t, reference.shape[4], reference.shape[2], reference.shape[3])
|
||||
resized = TVF.resize(
|
||||
reference_flat,
|
||||
size=(height, width),
|
||||
interpolation=InterpolationMode.BICUBIC,
|
||||
antialias=not (isinstance(reference_flat, torch.Tensor) and reference_flat.device.type == "mps"),
|
||||
)
|
||||
return resized.reshape(b, t, resized.shape[1], height, width).permute(0, 1, 3, 4, 2)
|
||||
|
||||
|
||||
class SeedVR2Conditioning(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="SeedVR2Conditioning",
|
||||
display_name="Apply SeedVR2 Conditioning",
|
||||
category="model/conditioning",
|
||||
description="Build SeedVR2 positive/negative conditioning from a VAE latent.",
|
||||
search_aliases=["seedvr2", "upscale", "conditioning"],
|
||||
inputs=[
|
||||
io.Model.Input("model", tooltip="The SeedVR2 model."),
|
||||
io.Latent.Input("vae_conditioning", display_name="latent"),
|
||||
],
|
||||
outputs=[
|
||||
io.Conditioning.Output(display_name="positive", tooltip="The positive conditioning for sampling."),
|
||||
io.Conditioning.Output(display_name="negative", tooltip="The negative conditioning for sampling."),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, model, vae_conditioning) -> io.NodeOutput:
|
||||
|
||||
vae_conditioning = vae_conditioning["samples"]
|
||||
if vae_conditioning.ndim != 5:
|
||||
raise ValueError(
|
||||
"SeedVR2Conditioning expects a 5-D VAE latent in Comfy "
|
||||
f"channel-first layout; got shape {tuple(vae_conditioning.shape)}."
|
||||
)
|
||||
if vae_conditioning.shape[1] != SEEDVR2_LATENT_CHANNELS:
|
||||
if vae_conditioning.shape[-1] == SEEDVR2_LATENT_CHANNELS:
|
||||
raise ValueError(
|
||||
"SeedVR2Conditioning expects SeedVR2 VAE latents in Comfy "
|
||||
f"channel-first layout (B, {SEEDVR2_LATENT_CHANNELS}, T, H, W); "
|
||||
f"got channel-last shape {tuple(vae_conditioning.shape)}."
|
||||
)
|
||||
raise ValueError(
|
||||
"SeedVR2Conditioning expects SeedVR2 VAE latents with "
|
||||
f"{SEEDVR2_LATENT_CHANNELS} channels; got shape {tuple(vae_conditioning.shape)}."
|
||||
)
|
||||
vae_conditioning = vae_conditioning.movedim(1, -1).contiguous()
|
||||
model = _resolve_seedvr2_diffusion_model(model)
|
||||
pos_cond = model.positive_conditioning
|
||||
neg_cond = model.negative_conditioning
|
||||
|
||||
mask = vae_conditioning.new_ones(vae_conditioning.shape[:-1] + (1,))
|
||||
condition = torch.cat((vae_conditioning, mask), dim=-1)
|
||||
condition = condition.movedim(-1, 1)
|
||||
|
||||
negative = [[neg_cond.unsqueeze(0), {"condition": condition}]]
|
||||
positive = [[pos_cond.unsqueeze(0), {"condition": condition}]]
|
||||
|
||||
return io.NodeOutput(positive, negative)
|
||||
|
||||
def _seedvr2_chunk_crossfade_weights(overlap, device, dtype):
|
||||
"""Descending previous-chunk weights across the overlap (next chunk gets ``1 - w``): a Hann fade over the middle third, flat shoulders on the outer thirds."""
|
||||
ramp = torch.linspace(0.0, 1.0, steps=overlap, device=device, dtype=dtype)
|
||||
ramp = ((ramp - 1.0 / 3.0) / (1.0 / 3.0)).clamp(0.0, 1.0)
|
||||
return 0.5 + 0.5 * torch.cos(torch.pi * ramp)
|
||||
|
||||
|
||||
class SeedVR2TemporalChunk(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="SeedVR2TemporalChunk",
|
||||
display_name="Split SeedVR2 Latent",
|
||||
category="model/latent/batch",
|
||||
description="Split a SeedVR2 video latent into overlapping temporal chunks small enough to sample one at a time within VRAM, wiring latents outputs to both Apply SeedVR2 Conditioning and the sampler latent input before recombining with Merge SeedVR2 Latents.",
|
||||
search_aliases=["seedvr2", "split", "chunk", "temporal", "video upscale", "rebatch"],
|
||||
inputs=[
|
||||
io.Latent.Input("latent", tooltip="The VAE-encoded SeedVR2 latent to split."),
|
||||
io.Int.Input("temporal_overlap", default=0, min=0, max=16384,
|
||||
tooltip="Latent frames shared between adjacent chunks and crossfaded at merge; 0 = no overlap."),
|
||||
io.DynamicCombo.Input("chunking_mode",
|
||||
tooltip="manual = use frames_per_chunk exactly; auto = predict the largest chunk that fits free VRAM.",
|
||||
options=[
|
||||
io.DynamicCombo.Option("auto", []),
|
||||
io.DynamicCombo.Option("manual", [
|
||||
io.Int.Input("frames_per_chunk", default=21, min=1, max=16384, step=4,
|
||||
tooltip="Pixel frames per temporal chunk (4n+1: 1, 5, 9, 13, ...)."),
|
||||
]),
|
||||
]),
|
||||
],
|
||||
outputs=[
|
||||
io.Latent.Output(display_name="latents", is_output_list=True,
|
||||
tooltip="The temporal chunks in sequence order."),
|
||||
io.Int.Output(display_name="temporal_overlap",
|
||||
tooltip="The effective latent-frame overlap between adjacent chunks, for Merge SeedVR2 Latents."),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, latent, temporal_overlap, chunking_mode) -> io.NodeOutput:
|
||||
samples = latent["samples"]
|
||||
if samples.ndim != 5:
|
||||
raise ValueError(
|
||||
f"SeedVR2TemporalChunk: expected a 5-D video latent (B, C, T, H, W); "
|
||||
f"got shape {tuple(samples.shape)}."
|
||||
)
|
||||
if samples.shape[1] != SEEDVR2_LATENT_CHANNELS:
|
||||
raise ValueError(
|
||||
f"SeedVR2TemporalChunk: expected {SEEDVR2_LATENT_CHANNELS} latent channels; "
|
||||
f"got shape {tuple(samples.shape)}."
|
||||
)
|
||||
if temporal_overlap < 0:
|
||||
raise ValueError(
|
||||
f"SeedVR2TemporalChunk: temporal_overlap must be >= 0; got {temporal_overlap}."
|
||||
)
|
||||
mode = chunking_mode["chunking_mode"]
|
||||
if mode not in ("auto", "manual"):
|
||||
raise ValueError(
|
||||
f"SeedVR2TemporalChunk: chunking_mode must be 'auto' or 'manual'; "
|
||||
f"got {mode!r}."
|
||||
)
|
||||
t_latent = samples.shape[2]
|
||||
t_pixel = 4 * (t_latent - 1) + 1
|
||||
|
||||
if mode == "auto":
|
||||
free_gb = comfy.model_management.get_free_memory(
|
||||
comfy.model_management.get_torch_device()) / (1024 ** 3)
|
||||
mpx_per_frame = (samples.shape[0] * samples.shape[3] * samples.shape[4]) * (BYTEDANCE_VAE_SPATIAL_DOWNSAMPLE ** 2) / 1e6
|
||||
budget_gb = free_gb - SEEDVR2_CHUNK_RESERVED_GIB - SEEDVR2_CHUNK_SIGMA_K * SEEDVR2_CHUNK_SIGMA_GIB
|
||||
chunk_latent_max = max(1, int(budget_gb / (SEEDVR2_CHUNK_GIB_PER_MPX_FRAME * mpx_per_frame)))
|
||||
frames_per_chunk = min(4 * (chunk_latent_max - 1) + 1, t_pixel)
|
||||
logging.info(
|
||||
"SeedVR2TemporalChunk auto: free=%.2fGiB, %.2fMpx -> frames_per_chunk=%d (t_pixel=%d).",
|
||||
free_gb, mpx_per_frame, frames_per_chunk, t_pixel,
|
||||
)
|
||||
else:
|
||||
frames_per_chunk = chunking_mode["frames_per_chunk"]
|
||||
if frames_per_chunk < 1 or (frames_per_chunk - 1) % 4 != 0:
|
||||
raise ValueError(
|
||||
f"SeedVR2TemporalChunk: frames_per_chunk must be a 4n+1 pixel-frame count "
|
||||
f"(1, 5, 9, 13, 17, 21, ...); got {frames_per_chunk}."
|
||||
)
|
||||
|
||||
if t_pixel <= frames_per_chunk:
|
||||
return io.NodeOutput([latent], 0)
|
||||
|
||||
chunk_latent = (frames_per_chunk - 1) // 4 + 1
|
||||
temporal_overlap = min(temporal_overlap, chunk_latent - 1)
|
||||
step = chunk_latent - temporal_overlap
|
||||
|
||||
chunks = []
|
||||
for start in range(0, t_latent, step):
|
||||
end = min(start + chunk_latent, t_latent)
|
||||
chunk = latent.copy()
|
||||
chunk["samples"] = samples[:, :, start:end].contiguous()
|
||||
chunks.append(chunk)
|
||||
if end >= t_latent:
|
||||
break
|
||||
return io.NodeOutput(chunks, temporal_overlap)
|
||||
|
||||
|
||||
class SeedVR2TemporalMerge(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="SeedVR2TemporalMerge",
|
||||
display_name="Merge SeedVR2 Latents",
|
||||
category="model/latent/batch",
|
||||
is_input_list=True,
|
||||
description="Recombine sampled SeedVR2 latent temporal chunks into one latent, crossfading each overlap with a Hann window sized by the temporal_overlap wired from Split SeedVR2 Latent.",
|
||||
search_aliases=["seedvr2", "merge", "temporal", "hann", "crossfade"],
|
||||
inputs=[
|
||||
io.Latent.Input("latents", tooltip="The sampled temporal chunks in sequence order."),
|
||||
io.Int.Input("temporal_overlap", default=0, min=0, max=16384, force_input=True,
|
||||
tooltip="The temporal_overlap output of Split SeedVR2 Latent. 0 = plain concatenation."),
|
||||
],
|
||||
outputs=[
|
||||
io.Latent.Output(display_name="latent", tooltip="The recombined full-length latent."),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, latents, temporal_overlap) -> io.NodeOutput:
|
||||
temporal_overlap = temporal_overlap[0]
|
||||
if temporal_overlap < 0:
|
||||
raise ValueError(
|
||||
f"SeedVR2TemporalMerge: temporal_overlap must be >= 0; got {temporal_overlap}."
|
||||
)
|
||||
chunks = [entry["samples"] for entry in latents]
|
||||
first = chunks[0]
|
||||
if first.ndim != 5:
|
||||
raise ValueError(
|
||||
f"SeedVR2TemporalMerge: expected 5-D video latents (B, C, T, H, W); "
|
||||
f"chunk 0 has shape {tuple(first.shape)}."
|
||||
)
|
||||
for i, chunk in enumerate(chunks[1:], start=1):
|
||||
if chunk.shape[:2] != first.shape[:2] or chunk.shape[3:] != first.shape[3:]:
|
||||
raise ValueError(
|
||||
f"SeedVR2TemporalMerge: chunk {i} shape {tuple(chunk.shape)} does not "
|
||||
f"match chunk 0 shape {tuple(first.shape)} outside the temporal axis."
|
||||
)
|
||||
if i < len(chunks) - 1 and chunk.shape[2] != first.shape[2]:
|
||||
raise ValueError(
|
||||
f"SeedVR2TemporalMerge: chunk {i} has {chunk.shape[2]} latent frames but "
|
||||
f"chunk 0 has {first.shape[2]}; only the final chunk may be shorter."
|
||||
)
|
||||
|
||||
out = latents[0].copy()
|
||||
out.pop("noise_mask", None)
|
||||
|
||||
if len(chunks) == 1:
|
||||
out["samples"] = first
|
||||
return io.NodeOutput(out)
|
||||
if temporal_overlap == 0:
|
||||
out["samples"] = torch.cat(chunks, dim=2)
|
||||
return io.NodeOutput(out)
|
||||
|
||||
chunk_latent = first.shape[2]
|
||||
step = chunk_latent - min(temporal_overlap, chunk_latent - 1)
|
||||
t_total = step * (len(chunks) - 1) + chunks[-1].shape[2]
|
||||
b, c, _, h, w = first.shape
|
||||
merged = torch.empty((b, c, t_total, h, w), device=first.device, dtype=first.dtype)
|
||||
|
||||
merged[:, :, :chunk_latent] = first
|
||||
filled = chunk_latent
|
||||
for i, chunk in enumerate(chunks[1:], start=1):
|
||||
start = i * step
|
||||
end = start + chunk.shape[2]
|
||||
# Crossfade width is bounded by the previous fill frontier and by a runt
|
||||
# final chunk shorter than the configured overlap.
|
||||
fade = min(filled - start, chunk.shape[2])
|
||||
if fade > 0:
|
||||
w_prev = _seedvr2_chunk_crossfade_weights(
|
||||
fade, chunk.device, chunk.dtype).view(1, 1, fade, 1, 1)
|
||||
merged[:, :, start:start + fade] = (
|
||||
merged[:, :, start:start + fade] * w_prev + chunk[:, :, :fade] * (1.0 - w_prev)
|
||||
)
|
||||
merged[:, :, start + fade:end] = chunk[:, :, fade:]
|
||||
else:
|
||||
merged[:, :, start:end] = chunk
|
||||
filled = end
|
||||
|
||||
out["samples"] = merged
|
||||
return io.NodeOutput(out)
|
||||
|
||||
|
||||
class SeedVRExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
||||
return [
|
||||
SeedVR2Conditioning,
|
||||
SeedVR2Preprocess,
|
||||
SeedVR2PostProcessing,
|
||||
SeedVR2TemporalChunk,
|
||||
SeedVR2TemporalMerge,
|
||||
]
|
||||
|
||||
async def comfy_entrypoint() -> SeedVRExtension:
|
||||
return SeedVRExtension()
|
||||
@@ -1,71 +0,0 @@
|
||||
import os
|
||||
import json
|
||||
from typing_extensions import override
|
||||
from comfy_api.latest import io, ComfyExtension, ui
|
||||
import folder_paths
|
||||
|
||||
|
||||
class SaveTextNode(io.ComfyNode):
|
||||
"""Save text content to .txt, .md, or .json."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="SaveText",
|
||||
search_aliases=["save text", "write text", "export text"],
|
||||
display_name="Save Text",
|
||||
category="text",
|
||||
description="Save text content to a file in the output directory.",
|
||||
inputs=[
|
||||
io.String.Input("text", force_input=True),
|
||||
io.String.Input("filename_prefix", default="ComfyUI"),
|
||||
io.Combo.Input("format", options=["txt", "md", "json"], default="txt"),
|
||||
],
|
||||
outputs=[io.String.Output(display_name="text")],
|
||||
is_output_node=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, text, filename_prefix, format):
|
||||
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(
|
||||
filename_prefix,
|
||||
folder_paths.get_output_directory(),
|
||||
1,
|
||||
1,
|
||||
)
|
||||
|
||||
file = f"{filename}_{counter:05}.{format}"
|
||||
filepath = os.path.join(full_output_folder, file)
|
||||
|
||||
if format == "json":
|
||||
# tries to pretty print otherwise saves normally
|
||||
try:
|
||||
data = json.loads(text)
|
||||
with open(filepath, "w", encoding="utf-8") as f:
|
||||
json.dump(data, f, indent=2, ensure_ascii=False)
|
||||
except json.JSONDecodeError:
|
||||
with open(filepath, "w", encoding="utf-8") as f:
|
||||
f.write(text)
|
||||
else:
|
||||
with open(filepath, "w", encoding="utf-8") as f:
|
||||
f.write(text)
|
||||
|
||||
return io.NodeOutput(
|
||||
text,
|
||||
ui={
|
||||
"text": (text,),
|
||||
"files": [
|
||||
ui.SavedResult(file, subfolder, io.FolderType.output)
|
||||
]
|
||||
}
|
||||
)
|
||||
|
||||
class TextExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
||||
return [
|
||||
SaveTextNode
|
||||
]
|
||||
|
||||
async def comfy_entrypoint() -> TextExtension:
|
||||
return TextExtension()
|
||||
@@ -81,7 +81,7 @@ class SaveVideo(io.ComfyNode):
|
||||
display_name="Save Video",
|
||||
category="video",
|
||||
essentials_category="Basics",
|
||||
description="Saves the input videos to your ComfyUI output directory.",
|
||||
description="Saves the input images to your ComfyUI output directory.",
|
||||
inputs=[
|
||||
io.Video.Input("video", tooltip="The video to save."),
|
||||
io.String.Input("filename_prefix", default="video/ComfyUI", tooltip="The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."),
|
||||
|
||||
+29
-224
@@ -29,16 +29,11 @@ from comfy_execution.caching import (
|
||||
HierarchicalCache,
|
||||
LRUCache,
|
||||
RAMPressureCache,
|
||||
RAM_CACHE_LARGE_INTERMEDIATE,
|
||||
)
|
||||
from comfy_execution.graph import (
|
||||
DependencyCycleError,
|
||||
DynamicPrompt,
|
||||
ExecutionBlocker,
|
||||
ExecutionFailureBlocker,
|
||||
ExecutionList,
|
||||
NodeInputError,
|
||||
NodeNotFoundError,
|
||||
get_input_info,
|
||||
)
|
||||
from comfy_execution.graph_utils import GraphBuilder, is_link
|
||||
@@ -55,16 +50,6 @@ class ExecutionResult(Enum):
|
||||
SUCCESS = 0
|
||||
FAILURE = 1
|
||||
PENDING = 2
|
||||
BLOCKED = 3
|
||||
|
||||
|
||||
NODE_FAILURE_POLICY_FAIL_FAST = "fail_fast"
|
||||
NODE_FAILURE_POLICY_CONTINUE_INDEPENDENT = "continue_independent"
|
||||
NODE_FAILURE_POLICIES = frozenset({
|
||||
NODE_FAILURE_POLICY_FAIL_FAST,
|
||||
NODE_FAILURE_POLICY_CONTINUE_INDEPENDENT,
|
||||
})
|
||||
NODE_FAILURE_POLICY_EXTRA_DATA_KEY = "_node_failure_policy"
|
||||
|
||||
class DuplicateNodeError(Exception):
|
||||
pass
|
||||
@@ -118,47 +103,6 @@ class CacheEntry(NamedTuple):
|
||||
outputs: list
|
||||
|
||||
|
||||
def _failure_blocker_state(value):
|
||||
if isinstance(value, ExecutionFailureBlocker):
|
||||
return True, False
|
||||
if isinstance(value, dict):
|
||||
values = value.values()
|
||||
elif isinstance(value, (list, tuple)):
|
||||
values = value
|
||||
else:
|
||||
return False, True
|
||||
|
||||
has_failure_blocker = False
|
||||
has_normal_value = False
|
||||
for child in values:
|
||||
child_failure, child_normal = _failure_blocker_state(child)
|
||||
has_failure_blocker = has_failure_blocker or child_failure
|
||||
has_normal_value = has_normal_value or child_normal
|
||||
if has_failure_blocker and has_normal_value:
|
||||
break
|
||||
return has_failure_blocker, has_normal_value
|
||||
|
||||
|
||||
def _tag_node_raised(ex):
|
||||
# Marks an exception as raised by node code (vs execution machinery). Best effort: an
|
||||
# exception class rejecting attribute assignment must not mask the original error.
|
||||
try:
|
||||
ex._node_raised = True
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _is_recoverable_node_failure(ex):
|
||||
if isinstance(ex, (
|
||||
comfy.model_management.InterruptProcessingException,
|
||||
DependencyCycleError,
|
||||
NodeInputError,
|
||||
NodeNotFoundError,
|
||||
)):
|
||||
return False
|
||||
return not comfy.model_management.is_oom(ex)
|
||||
|
||||
|
||||
class CacheType(Enum):
|
||||
CLASSIC = 0
|
||||
LRU = 1
|
||||
@@ -207,7 +151,7 @@ class CacheSet:
|
||||
}
|
||||
return result
|
||||
|
||||
SENSITIVE_EXTRA_DATA_KEYS = ("auth_token_comfy_org", "api_key_comfy_org", NODE_FAILURE_POLICY_EXTRA_DATA_KEY)
|
||||
SENSITIVE_EXTRA_DATA_KEYS = ("auth_token_comfy_org", "api_key_comfy_org")
|
||||
|
||||
def get_input_data(inputs, class_def, unique_id, execution_list=None, dynprompt=None, extra_data={}):
|
||||
is_v3 = issubclass(class_def, _ComfyNodeInternal)
|
||||
@@ -310,22 +254,16 @@ async def _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, f
|
||||
async def process_inputs(inputs, index=None, input_is_list=False):
|
||||
if allow_interrupt:
|
||||
nodes.before_node_execution()
|
||||
# Prefer a runtime-failure blocker over a plain user blocker so failure
|
||||
# taint and retry semantics are not masked by input ordering.
|
||||
blocked_value = None
|
||||
for k, v in inputs.items():
|
||||
values = v if input_is_list else (v,)
|
||||
for e in values:
|
||||
if isinstance(e, ExecutionBlocker):
|
||||
if blocked_value is None or isinstance(e, ExecutionFailureBlocker):
|
||||
blocked_value = e
|
||||
if isinstance(blocked_value, ExecutionFailureBlocker):
|
||||
break
|
||||
if isinstance(blocked_value, ExecutionFailureBlocker):
|
||||
break
|
||||
execution_block = None
|
||||
if blocked_value is not None:
|
||||
execution_block = execution_block_cb(blocked_value) if execution_block_cb else blocked_value
|
||||
for k, v in inputs.items():
|
||||
if input_is_list:
|
||||
for e in v:
|
||||
if isinstance(e, ExecutionBlocker):
|
||||
v = e
|
||||
break
|
||||
if isinstance(v, ExecutionBlocker):
|
||||
execution_block = execution_block_cb(v) if execution_block_cb else v
|
||||
break
|
||||
if execution_block is None:
|
||||
if pre_execute_cb is not None and index is not None:
|
||||
pre_execute_cb(index)
|
||||
@@ -351,11 +289,7 @@ async def _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, f
|
||||
if inspect.iscoroutinefunction(f):
|
||||
async def async_wrapper(f, prompt_id, unique_id, list_index, args):
|
||||
with CurrentNodeContext(prompt_id, unique_id, list_index):
|
||||
try:
|
||||
return await f(**args)
|
||||
except Exception as ex:
|
||||
_tag_node_raised(ex)
|
||||
raise
|
||||
return await f(**args)
|
||||
task = asyncio.create_task(async_wrapper(f, prompt_id, unique_id, index, args=inputs))
|
||||
# Give the task a chance to execute without yielding
|
||||
await asyncio.sleep(0)
|
||||
@@ -366,11 +300,7 @@ async def _async_map_node_over_list(prompt_id, unique_id, obj, input_data_all, f
|
||||
results.append(task)
|
||||
else:
|
||||
with CurrentNodeContext(prompt_id, unique_id, index):
|
||||
try:
|
||||
result = f(**inputs)
|
||||
except Exception as ex:
|
||||
_tag_node_raised(ex)
|
||||
raise
|
||||
result = f(**inputs)
|
||||
results.append(result)
|
||||
else:
|
||||
results.append(execution_block)
|
||||
@@ -517,16 +447,11 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
|
||||
execution_list.cache_update(unique_id, cached)
|
||||
return (ExecutionResult.SUCCESS, None, None)
|
||||
|
||||
continue_on_failure = extra_data.get(NODE_FAILURE_POLICY_EXTRA_DATA_KEY) == NODE_FAILURE_POLICY_CONTINUE_INDEPENDENT
|
||||
input_data_all = None
|
||||
failure_blocked_invocations = 0
|
||||
successful_invocations = 0
|
||||
resumed_subgraph = False
|
||||
try:
|
||||
if unique_id in pending_async_nodes:
|
||||
pending_results, failure_blocked_invocations, successful_invocations = pending_async_nodes[unique_id]
|
||||
results = []
|
||||
for r in pending_results:
|
||||
for r in pending_async_nodes[unique_id]:
|
||||
if isinstance(r, asyncio.Task):
|
||||
try:
|
||||
results.append(r.result())
|
||||
@@ -539,8 +464,7 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
|
||||
del pending_async_nodes[unique_id]
|
||||
output_data, output_ui, has_subgraph = get_output_from_returns(results, class_def)
|
||||
elif unique_id in pending_subgraph_results:
|
||||
cached_results, failure_blocked_invocations = pending_subgraph_results[unique_id]
|
||||
resumed_subgraph = True
|
||||
cached_results = pending_subgraph_results[unique_id]
|
||||
resolved_outputs = []
|
||||
for is_subgraph, result in cached_results:
|
||||
if not is_subgraph:
|
||||
@@ -593,9 +517,6 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
|
||||
return (ExecutionResult.PENDING, None, None)
|
||||
|
||||
def execution_block_cb(block):
|
||||
nonlocal failure_blocked_invocations
|
||||
if isinstance(block, ExecutionFailureBlocker):
|
||||
failure_blocked_invocations += 1
|
||||
if block.message is not None:
|
||||
mes = {
|
||||
"prompt_id": prompt_id,
|
||||
@@ -614,8 +535,6 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
|
||||
else:
|
||||
return block
|
||||
def pre_execute_cb(call_index):
|
||||
nonlocal successful_invocations
|
||||
successful_invocations += 1
|
||||
# TODO - How to handle this with async functions without contextvars (which requires Python 3.12)?
|
||||
GraphBuilder.set_default_prefix(unique_id, call_index, 0)
|
||||
|
||||
@@ -630,7 +549,7 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
|
||||
comfy_aimdo.model_vbar.vbars_reset_watermark_limits()
|
||||
|
||||
if has_pending_tasks:
|
||||
pending_async_nodes[unique_id] = (output_data, failure_blocked_invocations, successful_invocations)
|
||||
pending_async_nodes[unique_id] = output_data
|
||||
unblock = execution_list.add_external_block(unique_id)
|
||||
async def await_completion():
|
||||
tasks = [x for x in output_data if isinstance(x, asyncio.Task)]
|
||||
@@ -685,26 +604,14 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
|
||||
for node_id in new_output_ids:
|
||||
execution_list.add_node(node_id)
|
||||
execution_list.cache_link(node_id, unique_id)
|
||||
if continue_on_failure:
|
||||
# Wait for expanded output nodes before finishing, so a failure inside the expansion taints
|
||||
# this node before its outputs can be cached as reusable.
|
||||
execution_list.add_completion_link(node_id, unique_id)
|
||||
for link in new_output_links:
|
||||
execution_list.add_strong_link(link[0], link[1], unique_id)
|
||||
pending_subgraph_results[unique_id] = (cached_outputs, failure_blocked_invocations)
|
||||
pending_subgraph_results[unique_id] = cached_outputs
|
||||
return (ExecutionResult.PENDING, None, None)
|
||||
|
||||
cache_entry = CacheEntry(ui=ui_outputs.get(unique_id), outputs=output_data)
|
||||
if continue_on_failure:
|
||||
has_failure_blocker, has_normal_output = _failure_blocker_state(output_data)
|
||||
else:
|
||||
has_failure_blocker, has_normal_output = False, True
|
||||
failure_tainted = has_failure_blocker or failure_blocked_invocations > 0 or execution_list.is_failure_tainted(unique_id)
|
||||
if failure_tainted:
|
||||
execution_list.mark_failure_tainted(unique_id)
|
||||
execution_list.cache_update(unique_id, cache_entry, transient=failure_tainted)
|
||||
if not failure_tainted:
|
||||
await caches.outputs.set(unique_id, cache_entry)
|
||||
execution_list.cache_update(unique_id, cache_entry)
|
||||
await caches.outputs.set(unique_id, cache_entry)
|
||||
|
||||
except comfy.model_management.InterruptProcessingException as iex:
|
||||
logging.info("Processing interrupted")
|
||||
@@ -741,23 +648,11 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
|
||||
"exception_message": "{}\n{}".format(ex, tips),
|
||||
"exception_type": exception_type,
|
||||
"traceback": traceback.format_tb(tb),
|
||||
"current_inputs": input_data_formatted,
|
||||
"node_raised": getattr(ex, "_node_raised", False),
|
||||
"current_inputs": input_data_formatted
|
||||
}
|
||||
|
||||
return (ExecutionResult.FAILURE, error_details, ex)
|
||||
|
||||
if resumed_subgraph:
|
||||
fully_failure_blocked = has_failure_blocker and not has_normal_output
|
||||
else:
|
||||
fully_failure_blocked = successful_invocations == 0 and (
|
||||
failure_blocked_invocations > 0
|
||||
or (has_failure_blocker and not has_normal_output)
|
||||
)
|
||||
if fully_failure_blocked:
|
||||
get_progress_state().block_progress(unique_id)
|
||||
return (ExecutionResult.BLOCKED, None, None)
|
||||
|
||||
get_progress_state().finish_progress(unique_id)
|
||||
executed.add(unique_id)
|
||||
|
||||
@@ -774,7 +669,6 @@ class PromptExecutor:
|
||||
self.caches = CacheSet(cache_type=self.cache_type, cache_args=self.cache_args)
|
||||
self.status_messages = []
|
||||
self.success = True
|
||||
self.execution_summary = None
|
||||
|
||||
def add_message(self, event, data: dict, broadcast: bool):
|
||||
data = {
|
||||
@@ -813,21 +707,6 @@ class PromptExecutor:
|
||||
}
|
||||
self.add_message("execution_error", mes, broadcast=False)
|
||||
|
||||
def handle_node_execution_error(self, prompt_id, prompt, current_outputs, executed, error):
|
||||
node_id = error["node_id"]
|
||||
mes = {
|
||||
"prompt_id": prompt_id,
|
||||
"node_id": node_id,
|
||||
"node_type": prompt[node_id]["class_type"],
|
||||
"executed": list(executed),
|
||||
"exception_message": error["exception_message"],
|
||||
"exception_type": error["exception_type"],
|
||||
"traceback": error["traceback"],
|
||||
"current_inputs": error["current_inputs"],
|
||||
"current_outputs": list(current_outputs),
|
||||
}
|
||||
self.add_message("execution_node_error", mes, broadcast=False)
|
||||
|
||||
def _notify_prompt_lifecycle(self, event: str, prompt_id: str):
|
||||
if not _has_cache_providers():
|
||||
return
|
||||
@@ -848,8 +727,6 @@ class PromptExecutor:
|
||||
set_preview_method(extra_data.get("preview_method"))
|
||||
|
||||
nodes.interrupt_processing(False)
|
||||
self.success = True
|
||||
self.execution_summary = None
|
||||
|
||||
if "client_id" in extra_data:
|
||||
self.server.client_id = extra_data["client_id"]
|
||||
@@ -894,75 +771,35 @@ class PromptExecutor:
|
||||
executed = set()
|
||||
execution_list = ExecutionList(dynamic_prompt, self.caches.outputs)
|
||||
current_outputs = self.caches.outputs.all_node_ids()
|
||||
output_targets = set(execute_outputs)
|
||||
failed_node_ids = set()
|
||||
blocked_node_ids = set()
|
||||
blocked_output_node_ids = set()
|
||||
successful_output_node_ids = set()
|
||||
node_failures = []
|
||||
continue_independent = extra_data.get(
|
||||
NODE_FAILURE_POLICY_EXTRA_DATA_KEY,
|
||||
NODE_FAILURE_POLICY_FAIL_FAST,
|
||||
) == NODE_FAILURE_POLICY_CONTINUE_INDEPENDENT
|
||||
for node_id in list(execute_outputs):
|
||||
execution_list.add_node(node_id)
|
||||
|
||||
while not execution_list.is_empty():
|
||||
node_id, error, ex = await execution_list.stage_node_execution()
|
||||
if error is not None:
|
||||
self.success = False
|
||||
self.handle_execution_error(prompt_id, dynamic_prompt.original_prompt, current_outputs, executed, error, ex)
|
||||
break
|
||||
|
||||
assert node_id is not None, "Node ID should not be None at this point"
|
||||
result, error, ex = await execute(self.server, dynamic_prompt, self.caches, node_id, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes, ui_node_outputs)
|
||||
self.success = result != ExecutionResult.FAILURE
|
||||
if result == ExecutionResult.FAILURE:
|
||||
if continue_independent and error.get("node_raised") and _is_recoverable_node_failure(ex):
|
||||
self.handle_node_execution_error(prompt_id, dynamic_prompt.original_prompt, current_outputs, executed, error)
|
||||
real_node_id = error["node_id"]
|
||||
failed_node_ids.add(real_node_id)
|
||||
node_failures.append((error, ex))
|
||||
blocker = ExecutionFailureBlocker(real_node_id)
|
||||
class_type = dynamic_prompt.get_node(node_id)["class_type"]
|
||||
class_def = nodes.NODE_CLASS_MAPPINGS[class_type]
|
||||
cache_entry = CacheEntry(
|
||||
ui=None,
|
||||
outputs=[[blocker] for _ in class_def.RETURN_TYPES],
|
||||
)
|
||||
execution_list.cache_update(node_id, cache_entry, transient=True)
|
||||
execution_list.mark_failure_tainted(node_id)
|
||||
get_progress_state().error_progress(node_id)
|
||||
execution_list.complete_node_execution()
|
||||
else:
|
||||
self.success = False
|
||||
self.handle_execution_error(prompt_id, dynamic_prompt.original_prompt, current_outputs, executed, error, ex)
|
||||
break
|
||||
self.handle_execution_error(prompt_id, dynamic_prompt.original_prompt, current_outputs, executed, error, ex)
|
||||
break
|
||||
elif result == ExecutionResult.PENDING:
|
||||
execution_list.unstage_node_execution()
|
||||
elif result == ExecutionResult.BLOCKED:
|
||||
real_node_id = dynamic_prompt.get_real_node_id(node_id)
|
||||
blocked_node_ids.add(real_node_id)
|
||||
if node_id in output_targets:
|
||||
blocked_output_node_ids.add(real_node_id)
|
||||
execution_list.complete_node_execution()
|
||||
else: # result == ExecutionResult.SUCCESS:
|
||||
if node_id in output_targets:
|
||||
successful_output_node_ids.add(dynamic_prompt.get_real_node_id(node_id))
|
||||
execution_list.complete_node_execution()
|
||||
|
||||
if self.cache_type == CacheType.RAM_PRESSURE:
|
||||
ram_release_callback(ram_inactive_headroom)
|
||||
ram_shortfall = ram_headroom - psutil.virtual_memory().available
|
||||
if ram_shortfall > 0:
|
||||
freed = ram_release_callback(ram_headroom, free_active=True, min_entry_size=RAM_CACHE_LARGE_INTERMEDIATE)
|
||||
ram_shortfall -= freed
|
||||
if comfy.model_management.should_free_pins_for_ram_pressure(ram_shortfall):
|
||||
freed = comfy.model_management.free_pins(ram_shortfall + 512 * (1024 ** 2))
|
||||
if freed < ram_shortfall:
|
||||
if freed > 64 * (1024 ** 2):
|
||||
# AIMDO MEM_DECOMMIT can outrun psutil.available catching up.
|
||||
time.sleep(0.05)
|
||||
ram_release_callback(ram_headroom, free_active=True)
|
||||
freed = comfy.model_management.free_pins(ram_shortfall + 512 * (1024 ** 2))
|
||||
if freed < ram_shortfall:
|
||||
if freed > 64 * (1024 ** 2):
|
||||
# AIMDO MEM_DECOMMIT can outrun psutil.available catching up.
|
||||
time.sleep(0.05)
|
||||
ram_release_callback(ram_headroom, free_active=True)
|
||||
else:
|
||||
# Only execute when the while-loop ends without break
|
||||
# Send cached UI for intermediate output nodes that weren't executed
|
||||
@@ -975,36 +812,7 @@ class PromptExecutor:
|
||||
if cached is not None:
|
||||
display_node_id = dynamic_prompt.get_display_node_id(node_id)
|
||||
_send_cached_ui(self.server, node_id, display_node_id, cached, prompt_id, ui_node_outputs)
|
||||
|
||||
if node_failures:
|
||||
self.execution_summary = {
|
||||
"has_errors": True,
|
||||
"execution_error_count": len(node_failures),
|
||||
"failed_node_ids": sorted(failed_node_ids)[:100],
|
||||
"blocked_node_ids": sorted(blocked_node_ids)[:100],
|
||||
"blocked_output_node_ids": sorted(blocked_output_node_ids)[:100],
|
||||
"successful_output_node_ids": sorted(successful_output_node_ids)[:100],
|
||||
}
|
||||
if successful_output_node_ids:
|
||||
self.execution_summary["completion_status"] = "partial_success"
|
||||
self.add_message(
|
||||
"execution_success",
|
||||
{"prompt_id": prompt_id, **self.execution_summary},
|
||||
broadcast=False,
|
||||
)
|
||||
else:
|
||||
self.success = False
|
||||
last_error, last_ex = node_failures[-1]
|
||||
self.handle_execution_error(
|
||||
prompt_id,
|
||||
dynamic_prompt.original_prompt,
|
||||
current_outputs,
|
||||
executed,
|
||||
last_error,
|
||||
last_ex,
|
||||
)
|
||||
else:
|
||||
self.add_message("execution_success", { "prompt_id": prompt_id }, broadcast=False)
|
||||
self.add_message("execution_success", { "prompt_id": prompt_id }, broadcast=False)
|
||||
|
||||
ui_outputs = {}
|
||||
meta_outputs = {}
|
||||
@@ -1462,7 +1270,6 @@ class PromptQueue:
|
||||
status_str: Literal['success', 'error']
|
||||
completed: bool
|
||||
messages: List[str]
|
||||
execution_summary: Optional[dict] = None
|
||||
|
||||
def task_done(self, item_id, history_result,
|
||||
status: Optional['PromptQueue.ExecutionStatus'], process_item=None):
|
||||
@@ -1474,8 +1281,6 @@ class PromptQueue:
|
||||
status_dict: Optional[dict] = None
|
||||
if status is not None:
|
||||
status_dict = copy.deepcopy(status._asdict())
|
||||
if status_dict.get("execution_summary") is None:
|
||||
del status_dict["execution_summary"]
|
||||
|
||||
if process_item is not None:
|
||||
prompt = process_item(prompt)
|
||||
|
||||
@@ -366,8 +366,7 @@ def prompt_worker(q, server_instance):
|
||||
status=execution.PromptQueue.ExecutionStatus(
|
||||
status_str='success' if e.success else 'error',
|
||||
completed=e.success,
|
||||
messages=e.status_messages,
|
||||
execution_summary=e.execution_summary), process_item=remove_sensitive)
|
||||
messages=e.status_messages), process_item=remove_sensitive)
|
||||
if server_instance.client_id is not None:
|
||||
server_instance.send_sync("executing", {"node": None, "prompt_id": prompt_id}, server_instance.client_id)
|
||||
|
||||
|
||||
@@ -1709,7 +1709,6 @@ class PreviewImage(SaveImage):
|
||||
self.compress_level = 1
|
||||
|
||||
SEARCH_ALIASES = ["preview", "preview image", "show image", "view image", "display image", "image viewer"]
|
||||
DESCRIPTION = "Preview the images without saving them to the ComfyUI output directory."
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
@@ -2459,7 +2458,6 @@ async def init_builtin_extra_nodes():
|
||||
"nodes_camera_trajectory.py",
|
||||
"nodes_edit_model.py",
|
||||
"nodes_tcfg.py",
|
||||
"nodes_seedvr.py",
|
||||
"nodes_context_windows.py",
|
||||
"nodes_qwen.py",
|
||||
"nodes_boogu.py",
|
||||
@@ -2505,7 +2503,6 @@ async def init_builtin_extra_nodes():
|
||||
"nodes_triposplat.py",
|
||||
"nodes_depth_anything_3.py",
|
||||
"nodes_seed.py",
|
||||
"nodes_text.py",
|
||||
]
|
||||
|
||||
import_failed = []
|
||||
|
||||
+13
-31
@@ -7,18 +7,18 @@ components:
|
||||
description: Timestamp when the asset was created
|
||||
format: date-time
|
||||
type: string
|
||||
display_name:
|
||||
description: Display name of the asset. Mirrors name for backwards compatibility.
|
||||
nullable: true
|
||||
type: string
|
||||
file_path:
|
||||
description: Relative path in global-namespace-root form (e.g. "models/checkpoints/flux.safetensors")
|
||||
nullable: true
|
||||
type: string
|
||||
hash:
|
||||
description: Blake3 hash of the asset content.
|
||||
pattern: ^blake3:[a-f0-9]{64}$
|
||||
type: string
|
||||
loader_path:
|
||||
description: The value a loader consumes to load this asset. Null when no loader can resolve the file.
|
||||
nullable: true
|
||||
type: string
|
||||
display_name:
|
||||
description: Human-facing label for the asset. Not unique.
|
||||
nullable: true
|
||||
type: string
|
||||
id:
|
||||
description: Unique identifier for the asset
|
||||
format: uuid
|
||||
@@ -144,14 +144,6 @@ components:
|
||||
AssetUpdated:
|
||||
description: Response returned when an existing asset is successfully updated.
|
||||
properties:
|
||||
display_name:
|
||||
description: Display name of the asset. Mirrors name for backwards compatibility.
|
||||
nullable: true
|
||||
type: string
|
||||
file_path:
|
||||
description: Relative path in global-namespace-root form (e.g. "models/checkpoints/flux.safetensors")
|
||||
nullable: true
|
||||
type: string
|
||||
hash:
|
||||
description: Blake3 hash of the asset content.
|
||||
pattern: ^blake3:[a-f0-9]{64}$
|
||||
@@ -922,13 +914,6 @@ components:
|
||||
number:
|
||||
description: Priority number for the queue (lower numbers have higher priority)
|
||||
type: number
|
||||
node_failure_policy:
|
||||
default: fail_fast
|
||||
description: Controls whether a runtime node failure terminates the prompt or only blocks dependent nodes
|
||||
enum:
|
||||
- fail_fast
|
||||
- continue_independent
|
||||
type: string
|
||||
partial_execution_targets:
|
||||
description: List of node names to execute
|
||||
items:
|
||||
@@ -1651,7 +1636,7 @@ paths:
|
||||
format: uuid
|
||||
type: string
|
||||
tags:
|
||||
description: JSON-encoded array of freeform tag strings, e.g. '["models","checkpoint"]'. Common types include "models", "input", "output", and "temp", but any tag can be used in any order.
|
||||
description: JSON-encoded array of tag strings. For new byte uploads, include exactly one destination role (`input`, `output`, or `models`); `models` uploads also require exactly one `model_type:<folder_name>` tag. Extra tags are stored as labels and do not create path components.
|
||||
type: string
|
||||
user_metadata:
|
||||
description: Custom JSON metadata as a string
|
||||
@@ -1836,7 +1821,7 @@ paths:
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/AssetUpdated'
|
||||
$ref: '#/components/schemas/Asset'
|
||||
description: Asset updated successfully
|
||||
"400":
|
||||
content:
|
||||
@@ -2477,6 +2462,9 @@ paths:
|
||||
supports_preview_metadata:
|
||||
description: Whether the server supports preview metadata
|
||||
type: boolean
|
||||
supports_model_type_tags:
|
||||
description: Whether the server supports namespaced model type asset tags
|
||||
type: boolean
|
||||
type: object
|
||||
description: Success
|
||||
headers:
|
||||
@@ -3304,12 +3292,6 @@ paths:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: Invalid request parameters
|
||||
"401":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: Unauthorized - Authentication required
|
||||
"500":
|
||||
content:
|
||||
application/json:
|
||||
|
||||
+3
-3
@@ -1,6 +1,6 @@
|
||||
comfyui-frontend-package==1.45.20
|
||||
comfyui-workflow-templates==0.11.9
|
||||
comfyui-embedded-docs==0.5.8
|
||||
comfyui-workflow-templates==0.11.6
|
||||
comfyui-embedded-docs==0.5.7
|
||||
torch
|
||||
torchsde
|
||||
torchvision
|
||||
@@ -22,7 +22,7 @@ alembic
|
||||
SQLAlchemy>=2.0.0
|
||||
filelock
|
||||
av>=16.0.0
|
||||
comfy-kitchen==0.2.19
|
||||
comfy-kitchen==0.2.16
|
||||
comfy-aimdo==0.4.10
|
||||
requests
|
||||
simpleeval>=1.0.0
|
||||
|
||||
@@ -39,7 +39,6 @@ from comfy.deploy_environment import get_deploy_environment
|
||||
import comfy.utils
|
||||
import comfy.model_management
|
||||
from comfy_api import feature_flags
|
||||
from comfy.comfy_api_env import get_environment_overrides
|
||||
import node_helpers
|
||||
from comfyui_version import __version__
|
||||
from app.frontend_management import FrontendManager, parse_version
|
||||
@@ -728,11 +727,7 @@ class PromptServer():
|
||||
|
||||
@routes.get("/features")
|
||||
async def get_features(request):
|
||||
features = feature_flags.get_server_features()
|
||||
overrides = get_environment_overrides()
|
||||
if overrides:
|
||||
features.update(overrides)
|
||||
return web.json_response(features)
|
||||
return web.json_response(feature_flags.get_server_features())
|
||||
|
||||
@routes.get("/prompt")
|
||||
async def get_prompt(request):
|
||||
@@ -1097,19 +1092,6 @@ class PromptServer():
|
||||
if "partial_execution_targets" in json_data:
|
||||
partial_execution_targets = json_data["partial_execution_targets"]
|
||||
|
||||
node_failure_policy = json_data.get(
|
||||
"node_failure_policy",
|
||||
execution.NODE_FAILURE_POLICY_FAIL_FAST,
|
||||
)
|
||||
if not isinstance(node_failure_policy, str) or node_failure_policy not in execution.NODE_FAILURE_POLICIES:
|
||||
error = {
|
||||
"type": "invalid_node_failure_policy",
|
||||
"message": "node_failure_policy must be 'fail_fast' or 'continue_independent'",
|
||||
"details": f"Invalid node_failure_policy: {node_failure_policy!r}",
|
||||
"extra_info": {},
|
||||
}
|
||||
return web.json_response({"error": error, "node_errors": {}}, status=400)
|
||||
|
||||
self.node_replace_manager.apply_replacements(prompt)
|
||||
|
||||
valid = await execution.validate_prompt(prompt_id, prompt, partial_execution_targets)
|
||||
@@ -1117,10 +1099,6 @@ class PromptServer():
|
||||
if "extra_data" in json_data:
|
||||
extra_data = json_data["extra_data"]
|
||||
|
||||
extra_data.pop(execution.NODE_FAILURE_POLICY_EXTRA_DATA_KEY, None)
|
||||
if node_failure_policy == execution.NODE_FAILURE_POLICY_CONTINUE_INDEPENDENT:
|
||||
extra_data[execution.NODE_FAILURE_POLICY_EXTRA_DATA_KEY] = node_failure_policy
|
||||
|
||||
if "client_id" in json_data:
|
||||
extra_data["client_id"] = json_data["client_id"]
|
||||
|
||||
|
||||
@@ -24,28 +24,6 @@ def app(model_manager):
|
||||
app.add_routes(routes)
|
||||
return app
|
||||
|
||||
async def test_get_model_folders_includes_registered_extensions(aiohttp_client, app, tmp_path):
|
||||
"""Folders expose their registered extension set verbatim; an empty list
|
||||
means match-all (filter_files_extensions semantics)."""
|
||||
with patch('folder_paths.folder_names_and_paths', {
|
||||
'test_checkpoints': ([str(tmp_path)], {'.safetensors', '.ckpt'}),
|
||||
'test_configs': ([str(tmp_path)], ['.yaml']),
|
||||
'test_match_all': ([str(tmp_path)], set()),
|
||||
'configs': ([str(tmp_path)], ['.yaml']),
|
||||
}):
|
||||
client = await aiohttp_client(app)
|
||||
response = await client.get('/experiment/models')
|
||||
|
||||
assert response.status == 200
|
||||
folders = {f['name']: f for f in await response.json()}
|
||||
|
||||
assert 'configs' not in folders # blocklisted
|
||||
assert folders['test_checkpoints']['folders'] == [str(tmp_path)]
|
||||
assert folders['test_checkpoints']['extensions'] == ['.ckpt', '.safetensors']
|
||||
assert folders['test_configs']['extensions'] == ['.yaml']
|
||||
# Match-all registrations are exposed honestly, not substituted.
|
||||
assert folders['test_match_all']['extensions'] == []
|
||||
|
||||
async def test_get_model_preview_safetensors(aiohttp_client, app, tmp_path):
|
||||
img = Image.new('RGB', (100, 100), 'white')
|
||||
img_byte_arr = BytesIO()
|
||||
|
||||
@@ -3,7 +3,7 @@ import uuid
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
from helpers import assert_hash_fields_consistent
|
||||
from helpers import assert_hash_fields_consistent, get_asset_filename
|
||||
|
||||
|
||||
def test_list_assets_paging_and_sort(http: requests.Session, api_base: str, asset_factory, make_asset_bytes):
|
||||
@@ -306,6 +306,132 @@ def test_list_assets_invalid_query_rejected(http: requests.Session, api_base: st
|
||||
assert body["error"]["code"] == error_code
|
||||
|
||||
|
||||
def test_list_assets_display_name_emitted(http, api_base, asset_factory, make_asset_bytes):
|
||||
"""`display_name` is emitted for every populated asset in list responses,
|
||||
derived from the storage path (category prefix + hash-based stored filename)."""
|
||||
scope = f"lf-dispname-{uuid.uuid4().hex[:6]}"
|
||||
tags = ["models", "model_type:checkpoints", "unit-tests", scope]
|
||||
asset_factory("dn_a.safetensors", tags, {}, make_asset_bytes("dn_a", 700))
|
||||
asset_factory("dn_b.safetensors", tags, {}, make_asset_bytes("dn_b", 700))
|
||||
|
||||
r = http.get(
|
||||
api_base + "/api/assets",
|
||||
params={"include_tags": f"unit-tests,{scope}", "limit": "50"},
|
||||
timeout=120,
|
||||
)
|
||||
body = r.json()
|
||||
assert r.status_code == 200, body
|
||||
assert body["assets"], "expected at least one asset"
|
||||
for asset in body["assets"]:
|
||||
assert "display_name" in asset, "populated asset must emit display_name"
|
||||
expected = "checkpoints/" + get_asset_filename(asset["asset_hash"], ".safetensors")
|
||||
assert asset["display_name"] == expected
|
||||
|
||||
|
||||
def test_list_assets_hash_filter_exact_match(http, api_base, asset_factory, make_asset_bytes):
|
||||
"""`hash` filters to assets whose content hash matches exactly."""
|
||||
scope = f"lf-hash-{uuid.uuid4().hex[:6]}"
|
||||
tags = ["models", "model_type:checkpoints", "unit-tests", scope]
|
||||
a = asset_factory("hf_a.safetensors", tags, {}, make_asset_bytes("hf_a", 1024))
|
||||
b = asset_factory("hf_b.safetensors", tags, {}, make_asset_bytes("hf_b", 2048))
|
||||
|
||||
target = a["hash"]
|
||||
assert target and a["hash"] != b["hash"], "fixtures must have distinct content hashes"
|
||||
|
||||
r = http.get(
|
||||
api_base + "/api/assets",
|
||||
params={"hash": target, "limit": "50"},
|
||||
timeout=120,
|
||||
)
|
||||
body = r.json()
|
||||
assert r.status_code == 200, body
|
||||
names = [x["name"] for x in body["assets"]]
|
||||
assert names == [a["name"]]
|
||||
assert body["total"] == 1
|
||||
|
||||
|
||||
def test_list_assets_hash_filter_no_match(http, api_base, asset_factory, make_asset_bytes):
|
||||
"""A well-formed but unknown hash returns an empty page (200)."""
|
||||
scope = f"lf-hash-none-{uuid.uuid4().hex[:6]}"
|
||||
tags = ["models", "model_type:checkpoints", "unit-tests", scope]
|
||||
asset_factory("hn_a.safetensors", tags, {}, make_asset_bytes("hn_a", 800))
|
||||
|
||||
unknown = "blake3:" + ("0" * 64)
|
||||
r = http.get(
|
||||
api_base + "/api/assets",
|
||||
params={"hash": unknown, "limit": "50"},
|
||||
timeout=120,
|
||||
)
|
||||
body = r.json()
|
||||
assert r.status_code == 200, body
|
||||
assert body["assets"] == []
|
||||
assert body["total"] == 0
|
||||
|
||||
|
||||
def test_list_assets_hash_filter_normalizes_case_and_whitespace(
|
||||
http, api_base, asset_factory, make_asset_bytes
|
||||
):
|
||||
"""`hash` is trimmed and lowercased before matching, so an upper-cased,
|
||||
space-padded value still matches the stored lowercase hash."""
|
||||
scope = f"lf-hashnorm-{uuid.uuid4().hex[:6]}"
|
||||
tags = ["models", "model_type:checkpoints", "unit-tests", scope]
|
||||
a = asset_factory("hnorm_a.safetensors", tags, {}, make_asset_bytes("hnorm_a", 1024))
|
||||
|
||||
target = a["hash"]
|
||||
assert target == target.lower(), "stored hash is expected to be lowercase"
|
||||
messy = f" {target.upper()} "
|
||||
|
||||
r = http.get(
|
||||
api_base + "/api/assets",
|
||||
params={"hash": messy, "limit": "50"},
|
||||
timeout=120,
|
||||
)
|
||||
body = r.json()
|
||||
assert r.status_code == 200, body
|
||||
names = [x["name"] for x in body["assets"]]
|
||||
assert names == [a["name"]]
|
||||
assert body["total"] == 1
|
||||
|
||||
|
||||
def test_list_assets_hash_filter_empty_returns_empty_page(
|
||||
http, api_base, asset_factory, make_asset_bytes
|
||||
):
|
||||
"""An explicitly-supplied but empty `hash` (`?hash=`) is an exact-match miss
|
||||
and returns an empty page, rather than silently disabling the filter."""
|
||||
scope = f"lf-hashempty-{uuid.uuid4().hex[:6]}"
|
||||
tags = ["models", "model_type:checkpoints", "unit-tests", scope]
|
||||
asset_factory("he_a.safetensors", tags, {}, make_asset_bytes("he_a", 800))
|
||||
|
||||
r = http.get(
|
||||
api_base + "/api/assets",
|
||||
params={"hash": "", "limit": "50"},
|
||||
timeout=120,
|
||||
)
|
||||
body = r.json()
|
||||
assert r.status_code == 200, body
|
||||
assert body["assets"] == []
|
||||
assert body["total"] == 0
|
||||
|
||||
|
||||
def test_list_assets_include_public_accepted(http, api_base, asset_factory, make_asset_bytes):
|
||||
"""`include_public` is accepted for contract parity; core results are always
|
||||
the caller's own assets regardless of its value (the param is inert)."""
|
||||
scope = f"lf-incpub-{uuid.uuid4().hex[:6]}"
|
||||
tags = ["models", "model_type:checkpoints", "unit-tests", scope]
|
||||
a = asset_factory("ip_a.safetensors", tags, {}, make_asset_bytes("ip_a", 900))
|
||||
|
||||
for value in ("false", "true"):
|
||||
r = http.get(
|
||||
api_base + "/api/assets",
|
||||
params={"include_tags": f"unit-tests,{scope}", "include_public": value, "limit": "50"},
|
||||
timeout=120,
|
||||
)
|
||||
body = r.json()
|
||||
assert r.status_code == 200, body
|
||||
names = [x["name"] for x in body["assets"]]
|
||||
assert a["name"] in names, f"caller's own asset must be returned (include_public={value})"
|
||||
|
||||
|
||||
def test_list_assets_name_contains_literal_underscore(
|
||||
http,
|
||||
api_base,
|
||||
|
||||
@@ -1,186 +0,0 @@
|
||||
"""SeedVR2 conditioning node regression tests."""
|
||||
|
||||
import importlib
|
||||
import sys
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from comfy.cli_args import args as cli_args
|
||||
from comfy.ldm.seedvr.constants import SEEDVR2_LATENT_CHANNELS
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
cli_args.cpu = True
|
||||
|
||||
|
||||
_SENTINEL = object()
|
||||
_TARGETS = (
|
||||
("comfy.model_management", "comfy"),
|
||||
("comfy_extras.nodes_seedvr", "comfy_extras"),
|
||||
)
|
||||
|
||||
|
||||
def _import_nodes_seedvr_isolated():
|
||||
"""Import comfy_extras.nodes_seedvr with comfy.model_management mocked."""
|
||||
priors = []
|
||||
for mod_name, parent_name in _TARGETS:
|
||||
prior_mod = sys.modules.get(mod_name, _SENTINEL)
|
||||
parent = sys.modules.get(parent_name)
|
||||
attr = mod_name.split(".")[-1]
|
||||
prior_attr = (
|
||||
getattr(parent, attr, _SENTINEL) if parent is not None else _SENTINEL
|
||||
)
|
||||
priors.append((mod_name, parent_name, attr, prior_mod, prior_attr))
|
||||
|
||||
mock_mm = MagicMock()
|
||||
for fn in (
|
||||
"xformers_enabled", "xformers_enabled_vae",
|
||||
"pytorch_attention_enabled", "pytorch_attention_enabled_vae",
|
||||
"sage_attention_enabled", "flash_attention_enabled",
|
||||
"is_intel_xpu",
|
||||
):
|
||||
getattr(mock_mm, fn).return_value = False
|
||||
tv = torch.version.__version__.split(".")
|
||||
mock_mm.torch_version_numeric = (int(tv[0]), int(tv[1]))
|
||||
mock_mm.WINDOWS = False
|
||||
sys.modules["comfy.model_management"] = mock_mm
|
||||
if sys.modules.get("comfy") is None:
|
||||
importlib.import_module("comfy")
|
||||
comfy_pkg = sys.modules.get("comfy")
|
||||
if comfy_pkg is not None:
|
||||
setattr(comfy_pkg, "model_management", mock_mm)
|
||||
nodes_seedvr = sys.modules.get("comfy_extras.nodes_seedvr") or (
|
||||
importlib.import_module("comfy_extras.nodes_seedvr")
|
||||
)
|
||||
|
||||
def _restore():
|
||||
for mod_name, parent_name, attr, prior_mod, prior_attr in priors:
|
||||
if prior_mod is _SENTINEL:
|
||||
sys.modules.pop(mod_name, None)
|
||||
else:
|
||||
sys.modules[mod_name] = prior_mod
|
||||
parent = sys.modules.get(parent_name)
|
||||
if parent is None:
|
||||
continue
|
||||
if prior_attr is _SENTINEL:
|
||||
if hasattr(parent, attr):
|
||||
delattr(parent, attr)
|
||||
else:
|
||||
setattr(parent, attr, prior_attr)
|
||||
|
||||
return nodes_seedvr, _restore
|
||||
|
||||
|
||||
class _Rope(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.freqs = nn.Parameter(torch.zeros(4))
|
||||
|
||||
|
||||
class _Block(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.rope = _Rope()
|
||||
|
||||
|
||||
class _DiffusionModel(nn.Module):
|
||||
def __init__(self, n_blocks=3, conditioning_dtype=torch.float32):
|
||||
super().__init__()
|
||||
self.blocks = nn.ModuleList([_Block() for _ in range(n_blocks)])
|
||||
self.register_buffer("positive_conditioning", torch.ones((2, 4), dtype=conditioning_dtype))
|
||||
self.register_buffer("negative_conditioning", torch.zeros((3, 4), dtype=conditioning_dtype))
|
||||
|
||||
|
||||
class _ModelInner:
|
||||
def __init__(self, diffusion_model):
|
||||
self.diffusion_model = diffusion_model
|
||||
|
||||
|
||||
class _ModelPatcher:
|
||||
def __init__(self, diffusion_model):
|
||||
self.model = _ModelInner(diffusion_model)
|
||||
|
||||
|
||||
def test_seedvr2_conditioning_schema_exposes_conditioning_outputs():
|
||||
nodes_seedvr, restore = _import_nodes_seedvr_isolated()
|
||||
try:
|
||||
schema = nodes_seedvr.SeedVR2Conditioning.define_schema()
|
||||
assert [input_item.id for input_item in schema.inputs] == [
|
||||
"model",
|
||||
"vae_conditioning",
|
||||
]
|
||||
assert schema.inputs[1].display_name == "latent"
|
||||
assert [output.display_name for output in schema.outputs] == [
|
||||
"positive",
|
||||
"negative",
|
||||
]
|
||||
finally:
|
||||
restore()
|
||||
|
||||
|
||||
def test_seedvr2_conditioning_rejects_wrong_latent_channels():
|
||||
nodes_seedvr, restore = _import_nodes_seedvr_isolated()
|
||||
try:
|
||||
patcher = _ModelPatcher(_DiffusionModel())
|
||||
vae_conditioning = {"samples": torch.zeros(1, 8, 2, 2, 2)}
|
||||
|
||||
with pytest.raises(ValueError, match=f"{SEEDVR2_LATENT_CHANNELS} channels"):
|
||||
nodes_seedvr.SeedVR2Conditioning.execute(patcher, vae_conditioning)
|
||||
finally:
|
||||
restore()
|
||||
|
||||
|
||||
def test_seedvr2_conditioning_returns_conditioning_deterministically():
|
||||
nodes_seedvr, restore = _import_nodes_seedvr_isolated()
|
||||
try:
|
||||
diffusion_model = _DiffusionModel()
|
||||
patcher = _ModelPatcher(diffusion_model)
|
||||
samples = torch.arange(
|
||||
1,
|
||||
1 + SEEDVR2_LATENT_CHANNELS * 3 * 2 * 2,
|
||||
dtype=torch.float32,
|
||||
).reshape(1, SEEDVR2_LATENT_CHANNELS, 3, 2, 2)
|
||||
vae_conditioning = {"samples": samples}
|
||||
|
||||
first_positive, first_negative = (
|
||||
nodes_seedvr.SeedVR2Conditioning.execute(
|
||||
patcher,
|
||||
vae_conditioning,
|
||||
)
|
||||
)
|
||||
second_positive, second_negative = (
|
||||
nodes_seedvr.SeedVR2Conditioning.execute(
|
||||
patcher,
|
||||
vae_conditioning,
|
||||
)
|
||||
)
|
||||
|
||||
channel_last = samples.movedim(1, -1).contiguous()
|
||||
expected_condition = torch.cat(
|
||||
[
|
||||
channel_last,
|
||||
torch.ones((*channel_last.shape[:-1], 1)),
|
||||
],
|
||||
dim=-1,
|
||||
).movedim(-1, 1)
|
||||
|
||||
assert torch.equal(
|
||||
first_positive[0][1]["condition"],
|
||||
expected_condition,
|
||||
)
|
||||
assert torch.equal(
|
||||
second_positive[0][1]["condition"],
|
||||
expected_condition,
|
||||
)
|
||||
assert torch.equal(
|
||||
first_negative[0][1]["condition"],
|
||||
expected_condition,
|
||||
)
|
||||
assert torch.equal(
|
||||
second_negative[0][1]["condition"],
|
||||
expected_condition,
|
||||
)
|
||||
finally:
|
||||
restore()
|
||||
@@ -1,55 +0,0 @@
|
||||
import importlib
|
||||
import inspect
|
||||
import sys
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import torch
|
||||
|
||||
from comfy.cli_args import args as cli_args
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
cli_args.cpu = True
|
||||
|
||||
|
||||
def test_seedvr_node_signature_matches_schema():
|
||||
mock_mm = MagicMock()
|
||||
mock_mm.xformers_enabled.return_value = False
|
||||
mock_mm.xformers_enabled_vae.return_value = False
|
||||
mock_mm.sage_attention_enabled.return_value = False
|
||||
mock_mm.flash_attention_enabled.return_value = False
|
||||
|
||||
sentinel = object()
|
||||
prior_cpu = cli_args.cpu
|
||||
cli_args.cpu = True
|
||||
prior_module = sys.modules.get("comfy_extras.nodes_seedvr", sentinel)
|
||||
comfy_pkg = sys.modules.get("comfy")
|
||||
prior_mm_attr = getattr(comfy_pkg, "model_management", sentinel) if comfy_pkg else sentinel
|
||||
|
||||
with patch.dict(sys.modules, {"comfy.model_management": mock_mm}):
|
||||
if comfy_pkg is not None:
|
||||
setattr(comfy_pkg, "model_management", mock_mm)
|
||||
sys.modules.pop("comfy_extras.nodes_seedvr", None)
|
||||
try:
|
||||
nodes_seedvr = importlib.import_module("comfy_extras.nodes_seedvr")
|
||||
for node_cls in (nodes_seedvr.SeedVR2Preprocess, nodes_seedvr.SeedVR2PostProcessing, nodes_seedvr.SeedVR2Conditioning):
|
||||
schema_ids = [i.id for i in node_cls.define_schema().inputs]
|
||||
exec_params = [
|
||||
p for p in inspect.signature(node_cls.execute).parameters.keys()
|
||||
if p != "cls"
|
||||
]
|
||||
assert schema_ids == exec_params, (
|
||||
f"{node_cls.__name__} schema/execute drift: "
|
||||
f"schema_ids={schema_ids}, exec_params={exec_params}"
|
||||
)
|
||||
finally:
|
||||
cli_args.cpu = prior_cpu
|
||||
if prior_module is sentinel:
|
||||
sys.modules.pop("comfy_extras.nodes_seedvr", None)
|
||||
else:
|
||||
sys.modules["comfy_extras.nodes_seedvr"] = prior_module
|
||||
if comfy_pkg is not None:
|
||||
if prior_mm_attr is sentinel:
|
||||
if hasattr(comfy_pkg, "model_management"):
|
||||
delattr(comfy_pkg, "model_management")
|
||||
else:
|
||||
setattr(comfy_pkg, "model_management", prior_mm_attr)
|
||||
@@ -1,51 +0,0 @@
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from comfy.cli_args import args as cli_args
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
cli_args.cpu = True
|
||||
|
||||
from comfy_extras import nodes_seedvr # noqa: E402
|
||||
|
||||
|
||||
def _schema_ids(items):
|
||||
return [item.id for item in items]
|
||||
|
||||
|
||||
def test_seedvr2_post_processing_schema():
|
||||
schema = nodes_seedvr.SeedVR2PostProcessing.define_schema()
|
||||
|
||||
assert _schema_ids(schema.inputs) == ["images", "original_resized_images", "color_correction_method"]
|
||||
assert schema.inputs[2].options == ["lab", "wavelet", "adain", "none"]
|
||||
assert schema.inputs[2].default == "lab"
|
||||
assert schema.outputs[0].get_io_type() == "IMAGE"
|
||||
|
||||
|
||||
def test_seedvr2_post_processing_oom_error_uses_color_correction_method(monkeypatch):
|
||||
decoded = torch.full((1, 3, 4, 4), 0.25)
|
||||
reference = torch.full((1, 3, 4, 4), 0.75)
|
||||
|
||||
def _lab(content, style):
|
||||
raise torch.cuda.OutOfMemoryError("CUDA out of memory")
|
||||
|
||||
monkeypatch.setattr(nodes_seedvr.comfy.model_management, "vae_device", lambda: torch.device("cpu"))
|
||||
monkeypatch.setattr(nodes_seedvr.comfy.model_management, "get_free_memory", lambda device: 1_000_000)
|
||||
|
||||
with patch.object(nodes_seedvr, "lab_color_transfer", _lab):
|
||||
with pytest.raises(RuntimeError) as excinfo:
|
||||
nodes_seedvr.SeedVR2PostProcessing._color_transfer_chunked(
|
||||
decoded, reference, torch.device("cpu"), "lab",
|
||||
)
|
||||
assert "color_correction_method=lab" in str(excinfo.value)
|
||||
assert " method=lab" not in str(excinfo.value)
|
||||
|
||||
|
||||
def test_seedvr2_post_processing_unknown_color_correction_method_raises():
|
||||
decoded = torch.zeros(1, 2, 4, 4, 3)
|
||||
original = torch.zeros(1, 2, 4, 4, 3)
|
||||
with pytest.raises(ValueError) as excinfo:
|
||||
nodes_seedvr.SeedVR2PostProcessing.execute(decoded, original, "bogus")
|
||||
assert "color_correction_method" in str(excinfo.value)
|
||||
@@ -1,77 +0,0 @@
|
||||
"""SeedVR2 temporal chunk/merge node regression tests."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from comfy.cli_args import args as cli_args
|
||||
from comfy.ldm.seedvr.constants import (
|
||||
BYTEDANCE_VAE_SPATIAL_DOWNSAMPLE,
|
||||
SEEDVR2_CHUNK_GIB_PER_MPX_FRAME,
|
||||
SEEDVR2_CHUNK_RESERVED_GIB,
|
||||
SEEDVR2_CHUNK_SIGMA_GIB,
|
||||
SEEDVR2_CHUNK_SIGMA_K,
|
||||
SEEDVR2_LATENT_CHANNELS,
|
||||
)
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
cli_args.cpu = True
|
||||
|
||||
import comfy.model_management # noqa: E402
|
||||
from comfy_extras.nodes_seedvr import SeedVR2TemporalChunk, SeedVR2TemporalMerge, _seedvr2_chunk_crossfade_weights # noqa: E402
|
||||
|
||||
def _latent(t_latent, h=8, w=8, b=1):
|
||||
g = torch.Generator().manual_seed(7)
|
||||
return {"samples": torch.randn(b, SEEDVR2_LATENT_CHANNELS, t_latent, h, w, generator=g)}
|
||||
|
||||
def _split(latent, frames_per_chunk, temporal_overlap, chunking_mode="manual"):
|
||||
combo = {"chunking_mode": chunking_mode}
|
||||
if chunking_mode != "auto":
|
||||
combo["frames_per_chunk"] = frames_per_chunk
|
||||
return SeedVR2TemporalChunk.execute(latent, temporal_overlap, combo).args
|
||||
|
||||
def _merge(chunks, temporal_overlap):
|
||||
return SeedVR2TemporalMerge.execute(chunks, [temporal_overlap]).args[0]
|
||||
|
||||
def test_chunk_temporal_windows_and_validation():
|
||||
with pytest.raises(ValueError, match="4n\\+1"):
|
||||
_split(_latent(9), 20, 0)
|
||||
with pytest.raises(ValueError, match="5-D"):
|
||||
_split({"samples": torch.zeros(1, SEEDVR2_LATENT_CHANNELS * 9, 8, 8)}, 21, 0)
|
||||
with pytest.raises(ValueError, match="chunking_mode"):
|
||||
_split(_latent(13), 21, 0, "adaptive")
|
||||
latent = _latent(13)
|
||||
chunks, overlap = _split(latent, 21, 2) # chunk_latent=6, step=4 -> [0:6], [4:10], [8:13]
|
||||
assert overlap == 2 and [c["samples"].shape[2] for c in chunks] == [6, 6, 5]
|
||||
assert all(torch.equal(c["samples"], latent["samples"][:, :, s:e]) for c, (s, e) in zip(chunks, [(0, 6), (4, 10), (8, 13)]))
|
||||
assert len(_split(_latent(13), 21, 999)[0]) == 8 # overlap clamps to chunk_latent-1 -> step=1
|
||||
assert (r := _split(_latent(5), 21, 3)) and len(r[0]) == 1 and r[1] == 0 # t_pixel <= 21: passthrough
|
||||
|
||||
def test_chunk_auto_mode_applies_vram_law(monkeypatch):
|
||||
mpx_per_frame = (32 * 32) * (BYTEDANCE_VAE_SPATIAL_DOWNSAMPLE ** 2) / 1e6
|
||||
free_gb = (
|
||||
SEEDVR2_CHUNK_RESERVED_GIB
|
||||
+ SEEDVR2_CHUNK_SIGMA_K * SEEDVR2_CHUNK_SIGMA_GIB
|
||||
+ 5.1 * SEEDVR2_CHUNK_GIB_PER_MPX_FRAME * mpx_per_frame
|
||||
)
|
||||
monkeypatch.setattr(comfy.model_management, "get_free_memory", lambda dev=None: free_gb * (1024 ** 3))
|
||||
assert [c["samples"].shape[2] for c in _split(_latent(13, h=32, w=32), 1, 0, "auto")[0]] == [5, 5, 3]
|
||||
assert _split(_latent(13, h=32, w=32, b=2), 1, 0, "auto")[0][0]["samples"].shape[2] == 2 # batch halves the chunk
|
||||
|
||||
def test_merge_crossfade_and_reassembly():
|
||||
latent = _latent(13)
|
||||
latent["noise_mask"] = torch.rand(1, 1, 13, 8, 8)
|
||||
latent["batch_index"] = [0]
|
||||
merged = _merge(_split(latent, 21, 0)[0], 0)
|
||||
assert torch.equal(merged["samples"], latent["samples"])
|
||||
assert "noise_mask" not in merged and merged["batch_index"] == [0]
|
||||
assert torch.allclose(_merge(_split(latent, 21, 3)[0], 3)["samples"], latent["samples"], atol=1e-6)
|
||||
w = _seedvr2_chunk_crossfade_weights(3, merged["samples"].device, merged["samples"].dtype)
|
||||
assert w[0] == 1.0 and w[-1] == 0.0 and torch.all(w[:-1] >= w[1:])
|
||||
ones, zeros = {"samples": torch.ones(1, SEEDVR2_LATENT_CHANNELS, 6, 8, 8)}, {"samples": torch.zeros(1, SEEDVR2_LATENT_CHANNELS, 6, 8, 8)}
|
||||
fused = _merge([ones, zeros], 3)["samples"] # overlap equals w: prev fades out, next fades in
|
||||
assert torch.equal(fused[:, :, 3:6], w.view(1, 1, 3, 1, 1).expand(1, SEEDVR2_LATENT_CHANNELS, 3, 8, 8))
|
||||
assert torch.equal(fused[:, :, :3], ones["samples"][:, :, :3]) and torch.equal(fused[:, :, 6:], zeros["samples"][:, :, :3])
|
||||
short = _split(latent, 21, 2)[0]
|
||||
short[0]["samples"] = short[0]["samples"][:, :, :4]
|
||||
with pytest.raises(ValueError, match="only the final chunk may be shorter"):
|
||||
_merge(short, 2)
|
||||
@@ -15,7 +15,7 @@ if not has_gpu():
|
||||
args.cpu = True
|
||||
|
||||
from comfy import ops
|
||||
from comfy.quant_ops import QUANT_ALGOS, QuantizedTensor
|
||||
from comfy.quant_ops import QuantizedTensor
|
||||
import comfy.utils
|
||||
|
||||
|
||||
@@ -283,59 +283,7 @@ class TestMixedPrecisionOps(unittest.TestCase):
|
||||
saved = model.state_dict()
|
||||
saved_conf = json.loads(saved["layer.comfy_quant"].numpy().tobytes())
|
||||
self.assertTrue(saved_conf["convrot"])
|
||||
|
||||
def test_convrot_w4a4_loads_into_params(self):
|
||||
"""ConvRot W4A4 checkpoints must load as the dedicated kitchen layout."""
|
||||
if "convrot_w4a4" not in QUANT_ALGOS:
|
||||
self.skipTest("comfy_kitchen does not provide ConvRot W4A4")
|
||||
|
||||
torch.manual_seed(456)
|
||||
layer_quant_config = {
|
||||
"layer": {
|
||||
"format": "convrot_w4a4",
|
||||
"convrot_groupsize": 256,
|
||||
"linear_dtype": "int8",
|
||||
}
|
||||
}
|
||||
weight = torch.randn(16, 256, dtype=torch.bfloat16)
|
||||
bias = torch.randn(16, dtype=torch.bfloat16)
|
||||
q_weight = QuantizedTensor.from_float(
|
||||
weight,
|
||||
"TensorCoreConvRotW4A4Layout",
|
||||
convrot_groupsize=256,
|
||||
quant_group_size=64,
|
||||
)
|
||||
state_dict = {
|
||||
"layer.weight": q_weight._qdata,
|
||||
"layer.bias": bias,
|
||||
"layer.weight_scale": q_weight._params.scale,
|
||||
}
|
||||
|
||||
state_dict, _ = comfy.utils.convert_old_quants(
|
||||
state_dict,
|
||||
metadata={"_quantization_metadata": json.dumps({"layers": layer_quant_config})},
|
||||
)
|
||||
model = torch.nn.Module()
|
||||
model.layer = ops.mixed_precision_ops({}).Linear(256, 16, device="cpu", dtype=torch.bfloat16)
|
||||
model.load_state_dict(state_dict, strict=False)
|
||||
|
||||
self.assertIsInstance(model.layer.weight, QuantizedTensor)
|
||||
self.assertEqual(model.layer.weight._layout_cls, "TensorCoreConvRotW4A4Layout")
|
||||
self.assertEqual(model.layer.weight._params.convrot_groupsize, 256)
|
||||
self.assertEqual(model.layer.weight._params.quant_group_size, 64)
|
||||
self.assertEqual(model.layer.weight._params.linear_dtype, "int8")
|
||||
|
||||
input_tensor = torch.randn(4, 256, dtype=torch.bfloat16)
|
||||
loaded_out = model.layer(input_tensor)
|
||||
ref_out = torch.nn.functional.linear(input_tensor, q_weight, bias)
|
||||
self.assertTrue(torch.equal(loaded_out, ref_out))
|
||||
|
||||
saved = model.state_dict()
|
||||
saved_conf = json.loads(saved["layer.comfy_quant"].numpy().tobytes())
|
||||
self.assertEqual(saved_conf["format"], "convrot_w4a4")
|
||||
self.assertEqual(saved_conf["convrot_groupsize"], 256)
|
||||
self.assertEqual(saved_conf["linear_dtype"], "int8")
|
||||
self.assertNotIn("quant_group_size", saved_conf)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -2,7 +2,7 @@ from collections import defaultdict
|
||||
|
||||
import torch
|
||||
|
||||
from comfy.model_detection import detect_unet_config, model_config_from_unet, model_config_from_unet_config
|
||||
from comfy.model_detection import detect_unet_config, model_config_from_unet_config
|
||||
import comfy.supported_models
|
||||
|
||||
|
||||
@@ -73,49 +73,6 @@ def _make_flux_schnell_comfyui_sd():
|
||||
return sd
|
||||
|
||||
|
||||
def _make_seedvr2_7b_separate_mm_sd():
|
||||
return {
|
||||
"blocks.35.mlp.vid.proj_out.weight": torch.empty(3072, 1),
|
||||
"positive_conditioning": torch.empty(58, 5120),
|
||||
"negative_conditioning": torch.empty(64, 5120),
|
||||
}
|
||||
|
||||
|
||||
def _make_seedvr2_7b_shared_mm_sd():
|
||||
return {
|
||||
"blocks.35.mlp.all.proj_in_gate.weight": torch.empty(1, 1),
|
||||
"positive_conditioning": torch.empty(58, 5120),
|
||||
"negative_conditioning": torch.empty(64, 5120),
|
||||
}
|
||||
|
||||
|
||||
def _make_seedvr2_3b_shared_mm_sd():
|
||||
return {
|
||||
"blocks.31.mlp.all.proj_in_gate.weight": torch.empty(1, 1),
|
||||
"positive_conditioning": torch.empty(58, 5120),
|
||||
"negative_conditioning": torch.empty(64, 5120),
|
||||
}
|
||||
|
||||
|
||||
def _make_pid_v1_5_sd(latent_proj_channels=16):
|
||||
sd = {
|
||||
"pixel_embedder.proj.weight": torch.empty(16, 3, device="meta"),
|
||||
"lq_proj.latent_proj.0.weight": torch.empty(1024, latent_proj_channels, 3, 3, device="meta"),
|
||||
"lq_proj.pit_head.weight": torch.empty(1536, 1024, device="meta"),
|
||||
"lq_proj.gate_modules.0.content_proj.weight": torch.empty(1, 3072, device="meta"),
|
||||
"pixel_blocks.0.attn.q_norm.weight": torch.empty(72, device="meta"),
|
||||
"pixel_blocks.0.adaLN_modulation.0.weight": torch.empty(24576, 1536, device="meta"),
|
||||
"pixel_blocks.0.adaLN_modulation.0.bias": torch.empty(24576, device="meta"),
|
||||
}
|
||||
for i in range(7):
|
||||
sd[f"lq_proj.gate_modules.{i}.log_alpha"] = torch.empty((), device="meta")
|
||||
return sd
|
||||
|
||||
|
||||
def _add_model_diffusion_prefix(sd):
|
||||
return {f"model.diffusion_model.{k}": v for k, v in sd.items()}
|
||||
|
||||
|
||||
class TestModelDetection:
|
||||
"""Verify that first-match model detection selects the correct model
|
||||
based on list ordering and unet_config specificity."""
|
||||
@@ -168,96 +125,6 @@ class TestModelDetection:
|
||||
assert model_config is not None
|
||||
assert type(model_config).__name__ == "FluxSchnell"
|
||||
|
||||
def test_seedvr2_7b_separate_mm_detection_config(self):
|
||||
sd = _make_seedvr2_7b_separate_mm_sd()
|
||||
unet_config = detect_unet_config(sd, "")
|
||||
|
||||
assert unet_config is not None
|
||||
assert unet_config["image_model"] == "seedvr2"
|
||||
assert unet_config["vid_dim"] == 3072
|
||||
assert unet_config["heads"] == 24
|
||||
assert unet_config["num_layers"] == 36
|
||||
assert unet_config["mm_layers"] == 36
|
||||
assert unet_config["mlp_type"] == "normal"
|
||||
assert unet_config["rope_type"] == "rope3d"
|
||||
assert unet_config["rope_dim"] == 64
|
||||
|
||||
def test_seedvr2_7b_shared_mm_detection_config(self):
|
||||
sd = _make_seedvr2_7b_shared_mm_sd()
|
||||
unet_config = detect_unet_config(sd, "")
|
||||
|
||||
assert unet_config is not None
|
||||
assert unet_config["image_model"] == "seedvr2"
|
||||
assert unet_config["vid_dim"] == 3072
|
||||
assert unet_config["heads"] == 24
|
||||
assert unet_config["num_layers"] == 36
|
||||
assert unet_config["mm_layers"] == 10
|
||||
assert unet_config["mlp_type"] == "swiglu"
|
||||
assert unet_config["rope_type"] == "rope3d"
|
||||
assert unet_config["rope_dim"] == 64
|
||||
|
||||
def test_seedvr2_3b_shared_mm_detection_config(self):
|
||||
sd = _make_seedvr2_3b_shared_mm_sd()
|
||||
unet_config = detect_unet_config(sd, "")
|
||||
|
||||
assert unet_config is not None
|
||||
assert unet_config["image_model"] == "seedvr2"
|
||||
assert unet_config["vid_dim"] == 2560
|
||||
assert unet_config["heads"] == 20
|
||||
assert unet_config["num_layers"] == 32
|
||||
assert unet_config["mlp_type"] == "swiglu"
|
||||
|
||||
def test_seedvr2_model_match_requires_conditioning_tensors(self):
|
||||
sd = _make_seedvr2_7b_shared_mm_sd()
|
||||
unet_config = detect_unet_config(sd, "")
|
||||
|
||||
assert type(model_config_from_unet_config(unet_config, sd)).__name__ == "SeedVR2"
|
||||
|
||||
del sd["positive_conditioning"]
|
||||
assert model_config_from_unet_config(unet_config, sd) is None
|
||||
|
||||
def test_seedvr2_model_match_accepts_full_checkpoint_prefix(self):
|
||||
sd = _add_model_diffusion_prefix(_make_seedvr2_7b_shared_mm_sd())
|
||||
|
||||
assert type(model_config_from_unet(sd, "model.diffusion_model.")).__name__ == "SeedVR2"
|
||||
|
||||
def test_pid_v1_5_detection(self):
|
||||
sd = _make_pid_v1_5_sd()
|
||||
unet_config = detect_unet_config(sd, "")
|
||||
|
||||
assert unet_config == {
|
||||
"image_model": "pid",
|
||||
"lq_latent_channels": 16,
|
||||
"lq_hidden_dim": 1024,
|
||||
"latent_spatial_down_factor": 8,
|
||||
"lq_interval": 2,
|
||||
"lq_latent_unpatchify_factor": 1,
|
||||
"lq_conv_padding_mode": "replicate",
|
||||
"lq_gate_per_token": True,
|
||||
"pit_lq_inject": True,
|
||||
"rope_ref_h": 2048,
|
||||
"rope_ref_w": 2048,
|
||||
}
|
||||
assert type(model_config_from_unet_config(unet_config, sd)).__name__ == "PiD"
|
||||
|
||||
def test_pid_v1_5_flux2_detection(self):
|
||||
unet_config = detect_unet_config(_make_pid_v1_5_sd(latent_proj_channels=32), "")
|
||||
|
||||
assert unet_config["lq_latent_channels"] == 128
|
||||
assert unet_config["latent_spatial_down_factor"] == 16
|
||||
assert unet_config["lq_latent_unpatchify_factor"] == 2
|
||||
|
||||
def test_pid_v1_5_pixel_adaln_conversion(self):
|
||||
sd = _make_pid_v1_5_sd()
|
||||
model_config = model_config_from_unet_config(detect_unet_config(sd, ""), sd)
|
||||
processed = model_config.process_unet_state_dict(sd)
|
||||
|
||||
assert processed["pixel_blocks.0.attn.q_norm.weight"].shape == (72,)
|
||||
assert processed["pixel_blocks.0.adaLN_modulation_msa.weight"].shape == (12288, 1536)
|
||||
assert processed["pixel_blocks.0.adaLN_modulation_mlp.weight"].shape == (12288, 1536)
|
||||
assert processed["pixel_blocks.0.adaLN_modulation_msa.bias"].shape == (12288,)
|
||||
assert processed["pixel_blocks.0.adaLN_modulation_mlp.bias"].shape == (12288,)
|
||||
|
||||
def test_unet_config_and_required_keys_combination_is_unique(self):
|
||||
"""Each model in the registry must have a unique combination of
|
||||
``unet_config`` and ``required_keys``. If two models share the same
|
||||
|
||||
@@ -1,74 +0,0 @@
|
||||
"""Regression tests for the SeedVR2 VAE forward return contract."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from comfy.cli_args import args as cli_args
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
cli_args.cpu = True
|
||||
|
||||
from comfy.ldm.seedvr.vae import SEEDVR2_LATENT_CHANNELS, VideoAutoencoderKL # noqa: E402
|
||||
|
||||
|
||||
_LATENT_SHAPE = (1, SEEDVR2_LATENT_CHANNELS, 2, 2, 2)
|
||||
_DECODED_SHAPE = (1, 3, 5, 16, 16)
|
||||
_INPUT_ENCODE_SHAPE = (1, 3, 5, 16, 16)
|
||||
_INPUT_DECODE_SHAPE = _LATENT_SHAPE
|
||||
|
||||
|
||||
class _StubVAE(VideoAutoencoderKL):
|
||||
def __init__(self):
|
||||
nn.Module.__init__(self)
|
||||
self._encode_out = torch.zeros(*_LATENT_SHAPE)
|
||||
self._decode_out = torch.zeros(*_DECODED_SHAPE)
|
||||
|
||||
def encode(self, x, return_dict=True):
|
||||
return self._encode_out
|
||||
|
||||
def decode_(self, z, return_dict=True):
|
||||
return self._decode_out
|
||||
|
||||
|
||||
def test_forward_encode_returns_tensor():
|
||||
vae = _StubVAE()
|
||||
x = torch.zeros(*_INPUT_ENCODE_SHAPE)
|
||||
result = vae.forward(x, mode="encode")
|
||||
assert type(result) is torch.Tensor
|
||||
assert result.shape == torch.Size(_LATENT_SHAPE)
|
||||
|
||||
|
||||
def test_forward_decode_returns_tensor():
|
||||
vae = _StubVAE()
|
||||
z = torch.zeros(*_INPUT_DECODE_SHAPE)
|
||||
result = vae.forward(z, mode="decode")
|
||||
assert type(result) is torch.Tensor
|
||||
assert result.shape == torch.Size(_DECODED_SHAPE)
|
||||
|
||||
|
||||
class _TupleReturningStubVAE(VideoAutoencoderKL):
|
||||
def __init__(self):
|
||||
nn.Module.__init__(self)
|
||||
self._encode_tensor = torch.zeros(*_LATENT_SHAPE)
|
||||
self._decode_tensor = torch.zeros(*_DECODED_SHAPE)
|
||||
|
||||
def encode(self, x, return_dict=True):
|
||||
return (self._encode_tensor,)
|
||||
|
||||
def decode_(self, z, return_dict=True):
|
||||
return (self._decode_tensor,)
|
||||
|
||||
|
||||
def test_forward_all_unwraps_one_tuple_at_each_step():
|
||||
vae = _TupleReturningStubVAE()
|
||||
x = torch.zeros(*_INPUT_ENCODE_SHAPE)
|
||||
result = vae.forward(x, mode="all")
|
||||
assert type(result) is torch.Tensor
|
||||
assert result.shape == torch.Size(_DECODED_SHAPE)
|
||||
|
||||
|
||||
def test_forward_rejects_unknown_mode():
|
||||
vae = _StubVAE()
|
||||
with pytest.raises(ValueError, match="Unknown SeedVR2 VAE forward mode"):
|
||||
vae.forward(torch.zeros(*_INPUT_ENCODE_SHAPE), mode="bogus")
|
||||
@@ -1,79 +0,0 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from comfy.cli_args import args as cli_args
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
cli_args.cpu = True
|
||||
|
||||
import comfy.sd
|
||||
import comfy.supported_models
|
||||
import comfy.ldm.seedvr.model as seedvr_model
|
||||
import comfy.ldm.seedvr.vae as seedvr_vae
|
||||
|
||||
|
||||
def test_seedvr2_fp16_manual_cast_only_for_bf16_device(monkeypatch):
|
||||
bf16_device = object()
|
||||
fp16_device = object()
|
||||
|
||||
monkeypatch.setattr(
|
||||
comfy.supported_models.comfy.model_management,
|
||||
"should_use_bf16",
|
||||
lambda device=None: device is bf16_device,
|
||||
)
|
||||
|
||||
bf16_config = comfy.supported_models.SeedVR2({"image_model": "seedvr2"})
|
||||
bf16_config.set_inference_dtype(torch.float16, None, device=bf16_device)
|
||||
assert bf16_config.manual_cast_dtype is torch.bfloat16
|
||||
|
||||
fp16_config = comfy.supported_models.SeedVR2({"image_model": "seedvr2"})
|
||||
fp16_config.set_inference_dtype(torch.float16, None, device=fp16_device)
|
||||
assert fp16_config.manual_cast_dtype is None
|
||||
|
||||
|
||||
def test_seedvr2_text_conditioning_accepts_cfg1_single_branch():
|
||||
context = torch.arange(6, dtype=torch.float32).reshape(1, 3, 2)
|
||||
|
||||
txt, txt_shape = seedvr_model.NaDiT._resolve_text_conditioning(object(), context, [0])
|
||||
|
||||
torch.testing.assert_close(txt, context.squeeze(0))
|
||||
torch.testing.assert_close(txt_shape, torch.tensor([[3]], device=context.device))
|
||||
|
||||
|
||||
def test_seedvr2_vae_decode_memory_covers_full_frame_lab_transfer():
|
||||
wrapper = seedvr_vae.VideoAutoencoderKLWrapper.__new__(seedvr_vae.VideoAutoencoderKLWrapper)
|
||||
latent_channels = seedvr_vae.SEEDVR2_LATENT_CHANNELS
|
||||
estimate = wrapper.comfy_memory_used_decode((1, latent_channels, 26, 120, 160))
|
||||
old_estimate = latent_channels * 120 * 160 * (4 * 8 * 8) * 2
|
||||
|
||||
assert estimate == 101 * 960 * 1280 * 160
|
||||
assert estimate > 15 * 1024 ** 3
|
||||
assert estimate > old_estimate * 100
|
||||
|
||||
|
||||
def test_seedvr2_vae_encode_preserves_compute_dtype(monkeypatch):
|
||||
wrapper = seedvr_vae.VideoAutoencoderKLWrapper.__new__(seedvr_vae.VideoAutoencoderKLWrapper)
|
||||
nn.Module.__init__(wrapper)
|
||||
wrapper._dummy = nn.Parameter(torch.empty(1, dtype=torch.float16))
|
||||
input_dtype = None
|
||||
|
||||
def encode(self, x):
|
||||
nonlocal input_dtype
|
||||
input_dtype = x.dtype
|
||||
return x
|
||||
|
||||
monkeypatch.setattr(seedvr_vae.VideoAutoencoderKL, "encode", encode)
|
||||
|
||||
x = torch.zeros((1, 3, 1, 8, 8), dtype=torch.float32)
|
||||
wrapper._encode_with_raw_latent(x)
|
||||
|
||||
assert input_dtype == torch.float32
|
||||
|
||||
|
||||
def test_seedvr2_vae_ops_cast_weights_to_compute_dtype():
|
||||
attention = seedvr_vae.Attention(query_dim=4, heads=1, dim_head=4).to(torch.float16)
|
||||
hidden_states = torch.zeros((1, 2, 4), dtype=torch.float32)
|
||||
|
||||
output = attention(hidden_states)
|
||||
|
||||
assert output.dtype == torch.float32
|
||||
@@ -1,169 +0,0 @@
|
||||
"""SeedVR2 internals regression tests."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from comfy.cli_args import args
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
args.cpu = True
|
||||
|
||||
import comfy.ldm.seedvr.model as seedvr_model # noqa: E402
|
||||
import comfy.ldm.seedvr.vae as vae_mod # noqa: E402
|
||||
import comfy.ldm.modules.attention as attention # noqa: E402
|
||||
import comfy.ops as comfy_ops # noqa: E402
|
||||
from comfy.ldm.seedvr.vae import ( # noqa: E402
|
||||
causal_norm_wrapper,
|
||||
set_norm_limit,
|
||||
)
|
||||
from comfy.ldm.seedvr.attention import var_attention_optimized_split # noqa: E402
|
||||
|
||||
|
||||
_NUM_CHANNELS = 8
|
||||
_NUM_GROUPS = 4
|
||||
_TENSOR_SHAPE = (1, 8, 2, 4, 4)
|
||||
|
||||
_GROUPNORM_SUBCLASSES = [
|
||||
pytest.param(comfy_ops.disable_weight_init.GroupNorm, id="disable_weight_init"),
|
||||
pytest.param(comfy_ops.manual_cast.GroupNorm, id="manual_cast"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("groupnorm_cls", _GROUPNORM_SUBCLASSES)
|
||||
def test_seedvr_groupnorm_low_limit_uses_chunked_groupnorm_path(groupnorm_cls):
|
||||
real_group_norm = vae_mod.F.group_norm
|
||||
set_norm_limit(1e-9)
|
||||
try:
|
||||
gn = groupnorm_cls(num_channels=_NUM_CHANNELS, num_groups=_NUM_GROUPS)
|
||||
gn.eval()
|
||||
|
||||
forward_hook_calls = []
|
||||
|
||||
def _hook(module, inputs, output):
|
||||
forward_hook_calls.append(tuple(inputs[0].shape))
|
||||
|
||||
spy_calls = []
|
||||
|
||||
def _group_norm_spy(input_tensor, num_groups_arg, *args, **kwargs):
|
||||
spy_calls.append({"num_groups": int(num_groups_arg)})
|
||||
return real_group_norm(input_tensor, num_groups_arg, *args, **kwargs)
|
||||
|
||||
handle = gn.register_forward_hook(_hook)
|
||||
try:
|
||||
with patch.object(vae_mod.F, "group_norm", side_effect=_group_norm_spy):
|
||||
out_tensor = causal_norm_wrapper(gn, torch.randn(*_TENSOR_SHAPE))
|
||||
finally:
|
||||
handle.remove()
|
||||
|
||||
full_calls = len(forward_hook_calls)
|
||||
chunked_calls = sum(1 for entry in spy_calls if entry["num_groups"] < _NUM_GROUPS)
|
||||
|
||||
assert tuple(int(s) for s in out_tensor.shape) == _TENSOR_SHAPE
|
||||
assert full_calls == 0, (
|
||||
f"low-limit GroupNorm gate must NOT take the full-forward path; got full_calls={full_calls}"
|
||||
)
|
||||
assert chunked_calls > 0, (
|
||||
f"low-limit GroupNorm gate must take the chunked path; got chunked_calls={chunked_calls}"
|
||||
)
|
||||
finally:
|
||||
set_norm_limit(None)
|
||||
|
||||
|
||||
def test_seedvr2_7b_swin_attention_forward_uses_optimized_var_attention(monkeypatch):
|
||||
dim = 8
|
||||
heads = 2
|
||||
head_dim = 4
|
||||
attn = seedvr_model.NaSwinAttention(
|
||||
vid_dim=dim,
|
||||
txt_dim=dim,
|
||||
heads=heads,
|
||||
head_dim=head_dim,
|
||||
qk_bias=False,
|
||||
qk_norm=comfy_ops.disable_weight_init.RMSNorm,
|
||||
qk_norm_eps=1e-6,
|
||||
rope_type=None,
|
||||
rope_dim=head_dim,
|
||||
shared_weights=False,
|
||||
window=(2, 1, 1),
|
||||
window_method="720pwin_by_size_bysize",
|
||||
version=True,
|
||||
device="cpu",
|
||||
dtype=torch.float32,
|
||||
operations=comfy_ops.disable_weight_init,
|
||||
)
|
||||
generator = torch.Generator(device="cpu").manual_seed(11)
|
||||
vid = torch.randn(8, dim, generator=generator)
|
||||
txt = torch.randn(3, dim, generator=generator)
|
||||
vid_shape = torch.tensor([[2, 2, 2]], dtype=torch.long)
|
||||
txt_shape = torch.tensor([[3]], dtype=torch.long)
|
||||
calls = []
|
||||
|
||||
def fake_optimized_var_attention(**kwargs):
|
||||
calls.append(kwargs)
|
||||
return kwargs["q"]
|
||||
|
||||
monkeypatch.setattr(seedvr_model, "optimized_var_attention", fake_optimized_var_attention)
|
||||
|
||||
vid_out, txt_out = attn(vid, txt, vid_shape, txt_shape, seedvr_model.Cache(disable=True))
|
||||
|
||||
assert tuple(vid_out.shape) == (8, dim)
|
||||
assert tuple(txt_out.shape) == (3, dim)
|
||||
assert len(calls) == 1
|
||||
call = calls[0]
|
||||
assert tuple(call["q"].shape) == (14, heads, head_dim)
|
||||
assert tuple(call["k"].shape) == (14, heads, head_dim)
|
||||
assert tuple(call["v"].shape) == (14, heads, head_dim)
|
||||
assert call["heads"] == heads
|
||||
assert call["skip_reshape"] is True
|
||||
assert call["skip_output_reshape"] is True
|
||||
assert call["cu_seqlens_q"] == [0, 7, 14]
|
||||
assert call["cu_seqlens_k"] == [0, 7, 14]
|
||||
|
||||
|
||||
def test_var_attention_optimized_split_calls_dense_backend_per_window(monkeypatch):
|
||||
heads = 2
|
||||
head_dim = 3
|
||||
q = torch.arange(30, dtype=torch.float32).reshape(5, heads, head_dim)
|
||||
k = q + 100
|
||||
v = q + 200
|
||||
cu = [0, 2, 5]
|
||||
calls = []
|
||||
|
||||
def fake_optimized_attention(q_arg, k_arg, v_arg, heads_arg, **kwargs):
|
||||
calls.append(
|
||||
{
|
||||
"q_shape": tuple(q_arg.shape),
|
||||
"k_shape": tuple(k_arg.shape),
|
||||
"v_shape": tuple(v_arg.shape),
|
||||
"heads": heads_arg,
|
||||
"kwargs": kwargs,
|
||||
}
|
||||
)
|
||||
return q_arg + v_arg
|
||||
|
||||
monkeypatch.setattr(attention, "optimized_attention", fake_optimized_attention)
|
||||
|
||||
out = var_attention_optimized_split(
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
heads,
|
||||
cu,
|
||||
cu,
|
||||
skip_reshape=True,
|
||||
skip_output_reshape=True,
|
||||
)
|
||||
|
||||
assert tuple(out.shape) == (5, heads, head_dim)
|
||||
assert len(calls) == 2
|
||||
assert calls[0]["q_shape"] == (1, heads, 2, head_dim)
|
||||
assert calls[1]["q_shape"] == (1, heads, 3, head_dim)
|
||||
assert all(call["heads"] == heads for call in calls)
|
||||
assert all(call["kwargs"]["skip_reshape"] is True for call in calls)
|
||||
assert all(call["kwargs"]["skip_output_reshape"] is True for call in calls)
|
||||
torch.testing.assert_close(out, q + v, rtol=0, atol=0)
|
||||
|
||||
@@ -1,320 +0,0 @@
|
||||
"""SeedVR2 model, latent-format, and VAE graph regression tests."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from comfy.cli_args import args
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
args.cpu = True
|
||||
|
||||
import comfy # noqa: E402
|
||||
import comfy.latent_formats # noqa: E402
|
||||
import comfy.ldm.seedvr.model as seedvr_model # noqa: E402
|
||||
import comfy.ldm.seedvr.vae as seedvr_vae_mod # noqa: E402
|
||||
import comfy.model_management # noqa: E402
|
||||
import comfy.ops as comfy_ops # noqa: E402
|
||||
import comfy.sample # noqa: E402
|
||||
import comfy.sd as sd_mod # noqa: E402
|
||||
import nodes as nodes_mod # noqa: E402
|
||||
from comfy.ldm.seedvr.model import NaDiT # noqa: E402
|
||||
|
||||
|
||||
_LATENT_CHANNELS = seedvr_vae_mod.SEEDVR2_LATENT_CHANNELS
|
||||
|
||||
|
||||
def _make_standin(positive_conditioning):
|
||||
class _StandIn(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.register_buffer(
|
||||
"positive_conditioning", positive_conditioning
|
||||
)
|
||||
|
||||
_resolve_text_conditioning = NaDiT._resolve_text_conditioning
|
||||
|
||||
return _StandIn()
|
||||
|
||||
|
||||
class _StubModule(nn.Module):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__()
|
||||
|
||||
|
||||
def _capture_last_layer_flags(monkeypatch, vid_dim: int, txt_in_dim: int) -> list[bool]:
|
||||
flags = []
|
||||
|
||||
class _Block(_StubModule):
|
||||
def __init__(self, *args, **kwargs):
|
||||
flags.append(kwargs["is_last_layer"])
|
||||
super().__init__()
|
||||
|
||||
monkeypatch.setattr(seedvr_model, "NaPatchIn", _StubModule)
|
||||
monkeypatch.setattr(seedvr_model, "NaPatchOut", _StubModule)
|
||||
monkeypatch.setattr(seedvr_model, "TimeEmbedding", _StubModule)
|
||||
monkeypatch.setattr(seedvr_model, "NaMMSRTransformerBlock", _Block)
|
||||
|
||||
seedvr_model.NaDiT(
|
||||
norm_eps=1e-5,
|
||||
num_layers=4,
|
||||
mlp_type="normal",
|
||||
vid_dim=vid_dim,
|
||||
txt_in_dim=txt_in_dim,
|
||||
heads=24,
|
||||
mm_layers=3,
|
||||
operations=comfy_ops.disable_weight_init,
|
||||
)
|
||||
|
||||
return flags
|
||||
|
||||
|
||||
class _Model:
|
||||
def __init__(self, latent_format):
|
||||
self._latent_format = latent_format
|
||||
|
||||
def get_model_object(self, name):
|
||||
assert name == "latent_format"
|
||||
return self._latent_format
|
||||
|
||||
|
||||
class _Patcher:
|
||||
def get_free_memory(self, device):
|
||||
return 1024 * 1024 * 1024
|
||||
|
||||
|
||||
class _EncodeWrapper(seedvr_vae_mod.VideoAutoencoderKLWrapper):
|
||||
def __init__(self, encoded):
|
||||
nn.Module.__init__(self)
|
||||
self.encoded = encoded
|
||||
self.spatial_downsample_factor = 8
|
||||
self.temporal_downsample_factor = 4
|
||||
self.seen = []
|
||||
|
||||
def encode(self, x):
|
||||
self.seen.append(tuple(x.shape))
|
||||
return self.encoded.to(device=x.device, dtype=x.dtype)
|
||||
|
||||
|
||||
class _DecodeWrapper(seedvr_vae_mod.VideoAutoencoderKLWrapper):
|
||||
def __init__(self):
|
||||
nn.Module.__init__(self)
|
||||
self.spatial_downsample_factor = 8
|
||||
self.temporal_downsample_factor = 4
|
||||
self.calls = []
|
||||
|
||||
def decode(self, z, seedvr2_tiling=None):
|
||||
self.calls.append({"shape": tuple(z.shape), "seedvr2_tiling": seedvr2_tiling})
|
||||
if z.ndim == 4:
|
||||
b, tc, h, w = z.shape
|
||||
t = tc // _LATENT_CHANNELS
|
||||
else:
|
||||
b, _, t, h, w = z.shape
|
||||
return torch.zeros(b, 3, t, h * 8, w * 8, dtype=z.dtype, device=z.device)
|
||||
|
||||
|
||||
def test_seedvr2_wrapper_public_encode_returns_tensor(monkeypatch):
|
||||
raw_latent = torch.full((1, _LATENT_CHANNELS, 1, 4, 5), 2.0)
|
||||
seen_shapes = []
|
||||
|
||||
def base_encode(self, x):
|
||||
seen_shapes.append(tuple(x.shape))
|
||||
return raw_latent.to(device=x.device, dtype=x.dtype)
|
||||
|
||||
monkeypatch.setattr(seedvr_vae_mod.VideoAutoencoderKL, "encode", base_encode)
|
||||
|
||||
vae = seedvr_vae_mod.VideoAutoencoderKLWrapper.__new__(seedvr_vae_mod.VideoAutoencoderKLWrapper)
|
||||
nn.Module.__init__(vae)
|
||||
vae._dummy = nn.Parameter(torch.zeros((), dtype=torch.float32))
|
||||
|
||||
latent = vae.encode(torch.zeros(1, 3, 32, 40))
|
||||
|
||||
assert type(latent) is torch.Tensor
|
||||
assert tuple(latent.shape) == (1, _LATENT_CHANNELS, 4, 5)
|
||||
assert seen_shapes == [(1, 3, 1, 32, 40)]
|
||||
|
||||
|
||||
def test_seedvr2_wrapper_private_encode_helper_keeps_raw_latent(monkeypatch):
|
||||
raw_latent = torch.full((1, _LATENT_CHANNELS, 1, 4, 5), 3.0)
|
||||
|
||||
def base_encode(self, x):
|
||||
return raw_latent.to(device=x.device, dtype=x.dtype)
|
||||
|
||||
monkeypatch.setattr(seedvr_vae_mod.VideoAutoencoderKL, "encode", base_encode)
|
||||
|
||||
vae = seedvr_vae_mod.VideoAutoencoderKLWrapper.__new__(seedvr_vae_mod.VideoAutoencoderKLWrapper)
|
||||
nn.Module.__init__(vae)
|
||||
vae._dummy = nn.Parameter(torch.zeros((), dtype=torch.float32))
|
||||
|
||||
latent, raw = vae._encode_with_raw_latent(torch.zeros(1, 3, 32, 40))
|
||||
|
||||
assert tuple(latent.shape) == (1, _LATENT_CHANNELS, 4, 5)
|
||||
assert tuple(raw.shape) == (1, _LATENT_CHANNELS, 1, 4, 5)
|
||||
assert torch.equal(raw, raw_latent)
|
||||
|
||||
|
||||
def _make_vae(wrapper):
|
||||
vae = sd_mod.VAE.__new__(sd_mod.VAE)
|
||||
vae.first_stage_model = wrapper
|
||||
vae.device = torch.device("cpu")
|
||||
vae.output_device = torch.device("cpu")
|
||||
vae.vae_dtype = torch.float32
|
||||
vae.latent_channels = _LATENT_CHANNELS
|
||||
vae.latent_dim = 3
|
||||
vae.downscale_ratio = (lambda a: max(0, (a + 3) // 4), 8, 8)
|
||||
vae.upscale_ratio = (lambda a: max(0, a * 4 - 3), 8, 8)
|
||||
vae.output_channels = 3
|
||||
vae.disable_offload = True
|
||||
vae.extra_1d_channel = None
|
||||
vae.crop_input = False
|
||||
vae.not_video = False
|
||||
vae.handles_tiling = isinstance(wrapper, seedvr_vae_mod.VideoAutoencoderKLWrapper)
|
||||
vae.format_encoded = wrapper.comfy_format_encoded
|
||||
vae.patcher = _Patcher()
|
||||
vae.process_input = lambda image: image
|
||||
vae.process_output = lambda image: image.add(1.0).div(2.0).clamp(0.0, 1.0)
|
||||
vae.vae_output_dtype = lambda: torch.float32
|
||||
vae.memory_used_encode = lambda shape, dtype: 1
|
||||
vae.memory_used_decode = lambda shape, dtype: 1
|
||||
vae.throw_exception_if_invalid = lambda: None
|
||||
vae.vae_encode_crop_pixels = lambda pixels: pixels
|
||||
vae.spacial_compression_decode = lambda: 8
|
||||
vae.temporal_compression_decode = lambda: 4
|
||||
return vae
|
||||
|
||||
|
||||
def test_missing_context_falls_back_to_positive_buffer():
|
||||
pos_buffer = torch.full((58, 5120), 7.0)
|
||||
standin = _make_standin(pos_buffer)
|
||||
txt, txt_shape = standin._resolve_text_conditioning(None)
|
||||
assert txt.shape == (58, 5120)
|
||||
assert (txt == 7.0).all(), (
|
||||
"fallback path must use the positive_conditioning buffer "
|
||||
"verbatim, not a zero tensor"
|
||||
)
|
||||
assert txt_shape.shape == (1, 1)
|
||||
assert txt_shape[0, 0].item() == 58
|
||||
|
||||
|
||||
def test_seedvr2_7b_keeps_final_block_text_path(monkeypatch):
|
||||
assert _capture_last_layer_flags(monkeypatch, vid_dim=3072, txt_in_dim=3072) == [
|
||||
False,
|
||||
False,
|
||||
False,
|
||||
False,
|
||||
]
|
||||
|
||||
|
||||
def test_seedvr2_7b_rope3d_matches_wrapper_oracle():
|
||||
rope = seedvr_model.get_na_rope("rope3d", dim=64)
|
||||
generator = torch.Generator(device="cpu").manual_seed(0)
|
||||
q = torch.randn(4, 2, 128, generator=generator)
|
||||
k = torch.randn(4, 2, 128, generator=generator)
|
||||
shape = torch.tensor([[1, 2, 2]], dtype=torch.long)
|
||||
freqs = rope.get_axial_freqs(1, 2, 2).reshape(4, -1)
|
||||
|
||||
expected_q = seedvr_model._apply_seedvr2_rotary_emb(
|
||||
freqs,
|
||||
q.permute(1, 0, 2).float(),
|
||||
).to(q.dtype).permute(1, 0, 2)
|
||||
expected_k = seedvr_model._apply_seedvr2_rotary_emb(
|
||||
freqs,
|
||||
k.permute(1, 0, 2).float(),
|
||||
).to(k.dtype).permute(1, 0, 2)
|
||||
|
||||
actual_q, actual_k = rope(q.clone(), k.clone(), shape, seedvr_model.Cache(disable=True))
|
||||
|
||||
torch.testing.assert_close(actual_q, expected_q, rtol=0, atol=0)
|
||||
torch.testing.assert_close(actual_k, expected_k, rtol=0, atol=0)
|
||||
|
||||
|
||||
def test_seedvr2_forward_requires_conditioning_latents():
|
||||
model = NaDiT.__new__(NaDiT)
|
||||
x = torch.zeros(1, _LATENT_CHANNELS, 1, 4, 5)
|
||||
|
||||
with pytest.raises(ValueError, match="requires conditioning latents"):
|
||||
NaDiT.forward(model, x, timestep=torch.tensor([1.0]), context=None)
|
||||
|
||||
|
||||
def test_seedvr2_latent_format_uses_native_video_latent_shape():
|
||||
latent_format = comfy.latent_formats.SeedVR2()
|
||||
latent_image = torch.zeros(1, 1, 4, 5)
|
||||
|
||||
fixed = comfy.sample.fix_empty_latent_channels(_Model(latent_format), latent_image)
|
||||
|
||||
assert latent_format.latent_channels == _LATENT_CHANNELS
|
||||
assert latent_format.latent_dimensions == 3
|
||||
assert fixed.shape == (1, _LATENT_CHANNELS, 1, 4, 5)
|
||||
|
||||
|
||||
def test_seedvr2_model_requires_native_5d_latent():
|
||||
latent = torch.zeros(1, _LATENT_CHANNELS, 2, 4, 5)
|
||||
assert NaDiT._check_seedvr2_video_latent(latent, _LATENT_CHANNELS, "latent") is latent
|
||||
|
||||
with pytest.raises(ValueError, match="5-D native latent"):
|
||||
NaDiT._check_seedvr2_video_latent(torch.zeros(1, _LATENT_CHANNELS * 2, 4, 5), _LATENT_CHANNELS, "latent")
|
||||
|
||||
|
||||
def test_seedvr2_encode_and_encode_tiled_preserve_native_latent_contract(monkeypatch):
|
||||
monkeypatch.setattr(sd_mod.model_management, "load_models_gpu", lambda *a, **k: None)
|
||||
|
||||
encoded = torch.full((1, _LATENT_CHANNELS, 2, 4, 5), 2.0)
|
||||
vae = _make_vae(_EncodeWrapper(encoded))
|
||||
pixels = torch.zeros(1, 5, 32, 40, 3)
|
||||
|
||||
node_output = nodes_mod.VAEEncode().encode(vae, pixels)[0]
|
||||
node_latent = node_output["samples"]
|
||||
assert set(node_output) == {"samples"}
|
||||
assert tuple(node_latent.shape) == (1, _LATENT_CHANNELS, 2, 4, 5)
|
||||
assert node_latent.dtype == torch.float32
|
||||
assert node_latent.stride()[-1] == 1
|
||||
assert torch.equal(node_latent, torch.full_like(node_latent, 2.0 * seedvr_vae_mod.BYTEDANCE_VAE_SCALING_FACTOR))
|
||||
|
||||
tiled = torch.full((1, _LATENT_CHANNELS, 2, 4, 5), 3.0)
|
||||
monkeypatch.setattr(seedvr_vae_mod, "tiled_vae", MagicMock(return_value=tiled))
|
||||
tiled_output = nodes_mod.VAEEncodeTiled().encode(
|
||||
vae,
|
||||
pixels,
|
||||
tile_size=512,
|
||||
overlap=64,
|
||||
temporal_size=16,
|
||||
temporal_overlap=4,
|
||||
)[0]
|
||||
tiled_latent = tiled_output["samples"]
|
||||
assert set(tiled_output) == {"samples"}
|
||||
assert tuple(tiled_latent.shape) == (1, _LATENT_CHANNELS, 2, 4, 5)
|
||||
assert tiled_latent.dtype == torch.float32
|
||||
assert torch.equal(tiled_latent, torch.full_like(tiled_latent, 3.0 * seedvr_vae_mod.BYTEDANCE_VAE_SCALING_FACTOR))
|
||||
|
||||
|
||||
def test_vaedecode_tiled_spatial_applies_temporal_discarded(monkeypatch):
|
||||
monkeypatch.setattr(sd_mod.model_management, "load_models_gpu", lambda *a, **k: None)
|
||||
vae = _make_vae(_DecodeWrapper())
|
||||
|
||||
nodes_mod.VAEDecodeTiled().decode(
|
||||
vae,
|
||||
{"samples": torch.zeros(1, _LATENT_CHANNELS, 2, 4, 5)},
|
||||
tile_size=512,
|
||||
overlap=64,
|
||||
temporal_size=16,
|
||||
temporal_overlap=4,
|
||||
)
|
||||
|
||||
# Spatial inputs flow through; temporal inputs are discarded as public tiling
|
||||
# knobs, but SeedVR2's internal MemoryState causal slicing is left intact.
|
||||
assert vae.first_stage_model.calls == [
|
||||
{
|
||||
"shape": (1, _LATENT_CHANNELS, 2, 4, 5),
|
||||
"seedvr2_tiling": {
|
||||
"enable_tiling": True,
|
||||
"tile_size": (512, 512),
|
||||
"tile_overlap": (64, 64),
|
||||
"temporal_size": None,
|
||||
"temporal_overlap": None,
|
||||
},
|
||||
}
|
||||
]
|
||||
@@ -1,94 +0,0 @@
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from comfy.cli_args import args as cli_args
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
cli_args.cpu = True
|
||||
|
||||
import comfy.ldm.seedvr.vae as vae_mod # noqa: E402
|
||||
from comfy_extras import nodes_seedvr # noqa: E402
|
||||
|
||||
|
||||
_LATENT_CHANNELS = vae_mod.SEEDVR2_LATENT_CHANNELS
|
||||
|
||||
|
||||
def _make_wrapper() -> vae_mod.VideoAutoencoderKLWrapper:
|
||||
wrapper = vae_mod.VideoAutoencoderKLWrapper.__new__(
|
||||
vae_mod.VideoAutoencoderKLWrapper
|
||||
)
|
||||
nn.Module.__init__(wrapper)
|
||||
return wrapper
|
||||
|
||||
|
||||
def _fingerprint_decode_(self, z, return_dict=True):
|
||||
b = int(z.shape[0])
|
||||
t = int(z.shape[2])
|
||||
h = int(z.shape[3])
|
||||
w = int(z.shape[4])
|
||||
out = torch.empty(b, 3, t, h * 8, w * 8)
|
||||
for batch_idx in range(b):
|
||||
out[batch_idx].fill_(float(batch_idx + 1))
|
||||
return out
|
||||
|
||||
|
||||
def _decode_with_patches(wrapper, z):
|
||||
with patch.object(vae_mod.VideoAutoencoderKL, "decode_", _fingerprint_decode_):
|
||||
return wrapper.decode(z)
|
||||
|
||||
|
||||
def test_decode_b2_t3_multi_frame_batch_unchanged():
|
||||
wrapper = _make_wrapper()
|
||||
|
||||
out = _decode_with_patches(wrapper, torch.zeros(2, _LATENT_CHANNELS * 3, 2, 2))
|
||||
|
||||
assert tuple(out.shape) == (2, 3, 3, 16, 16)
|
||||
|
||||
|
||||
class _Wrapper(vae_mod.VideoAutoencoderKLWrapper):
|
||||
def __init__(self):
|
||||
nn.Module.__init__(self)
|
||||
self.calls = []
|
||||
|
||||
def parameters(self):
|
||||
return iter([torch.nn.Parameter(torch.zeros(()))])
|
||||
|
||||
def _decode_stub(self, latent):
|
||||
self.calls.append(tuple(latent.shape))
|
||||
return torch.zeros(latent.shape[0], 3, latent.shape[2], latent.shape[3] * 8, latent.shape[4] * 8)
|
||||
|
||||
|
||||
def test_seedvr2_wrapper_decode_accepts_5d_channel_first_latents_without_preprocessor_state():
|
||||
wrapper = _Wrapper()
|
||||
|
||||
with patch.object(vae_mod.VideoAutoencoderKL, "decode_", _decode_stub):
|
||||
out = wrapper.decode(torch.zeros(1, _LATENT_CHANNELS, 2, 4, 5))
|
||||
|
||||
assert tuple(out.shape) == (1, 3, 2, 32, 40)
|
||||
assert wrapper.calls == [(1, _LATENT_CHANNELS, 2, 4, 5)]
|
||||
|
||||
|
||||
def test_seedvr2_wrapper_decode_rejects_wrong_rank_latents():
|
||||
wrapper = _Wrapper()
|
||||
|
||||
with pytest.raises(RuntimeError, match=r"latent input must be 4-D collapsed .* or 5-D"):
|
||||
wrapper.decode(torch.zeros(1, _LATENT_CHANNELS, 4))
|
||||
|
||||
|
||||
def _t_padded(t_in: int) -> int:
|
||||
if t_in == 1:
|
||||
return 1
|
||||
if t_in <= 4:
|
||||
return 5
|
||||
if (t_in - 1) % 4 == 0:
|
||||
return t_in
|
||||
return t_in + (4 - ((t_in - 1) % 4))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("t_in", [1, 5, 9])
|
||||
def test_t_padded_matches_cut_videos(t_in):
|
||||
dummy = torch.zeros(1, t_in, 1, 1, 1)
|
||||
assert nodes_seedvr.cut_videos(dummy).shape[1] == _t_padded(t_in)
|
||||
@@ -1,407 +0,0 @@
|
||||
from contextlib import ExitStack
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from comfy.cli_args import args as cli_args
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
cli_args.cpu = True
|
||||
|
||||
import comfy.ldm.seedvr.vae as vae_mod # noqa: E402
|
||||
import comfy.ldm.seedvr.vae as seedvr_vae_mod # noqa: E402
|
||||
import comfy.sd as sd_mod # noqa: E402
|
||||
from comfy.ldm.seedvr.vae import MemoryState, tiled_vae # noqa: E402
|
||||
|
||||
|
||||
_LATENT_CHANNELS = seedvr_vae_mod.SEEDVR2_LATENT_CHANNELS
|
||||
|
||||
|
||||
def test_runtime_decode_zero_temporal_size_preserves_model_slicing():
|
||||
class StubVAEModel(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.slicing_latent_min_size = 2
|
||||
self.spatial_downsample_factor = 8
|
||||
self.temporal_downsample_factor = 4
|
||||
self.device = torch.device("cpu")
|
||||
self.use_slicing = True
|
||||
self._dummy = torch.nn.Parameter(torch.zeros(1, dtype=torch.float32))
|
||||
self.decode_min_sizes = []
|
||||
self.memory_states = []
|
||||
|
||||
def decode_(self, t_chunk):
|
||||
self.decode_min_sizes.append(self.slicing_latent_min_size)
|
||||
return vae_mod.VideoAutoencoderKL.slicing_decode(self, t_chunk)
|
||||
|
||||
def _decode(self, z, memory_state=MemoryState.DISABLED, memory_cache=None):
|
||||
self.memory_states.append(memory_state)
|
||||
b, c, d, h, w = z.shape
|
||||
return torch.zeros((b, 3, d, h * 8, w * 8), dtype=z.dtype)
|
||||
|
||||
vae = StubVAEModel()
|
||||
z = torch.zeros((1, _LATENT_CHANNELS, 5, 8, 8), dtype=torch.float32)
|
||||
|
||||
tiled_vae(
|
||||
z,
|
||||
vae,
|
||||
tile_size=(64, 64),
|
||||
tile_overlap=(0, 0),
|
||||
temporal_size=0,
|
||||
temporal_overlap=0,
|
||||
encode=False,
|
||||
)
|
||||
|
||||
assert vae.decode_min_sizes == [2]
|
||||
assert vae.memory_states == [MemoryState.INITIALIZING, MemoryState.ACTIVE]
|
||||
assert vae.slicing_latent_min_size == 2
|
||||
|
||||
|
||||
def test_zero_temporal_size_preserves_min_size_when_encode_raises():
|
||||
class RaisingVAEModel(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.slicing_sample_min_size = 4
|
||||
self.spatial_downsample_factor = 8
|
||||
self.temporal_downsample_factor = 4
|
||||
self.device = torch.device("cpu")
|
||||
self._dummy = torch.nn.Parameter(torch.zeros(1, dtype=torch.float32))
|
||||
|
||||
def encode(self, t_chunk):
|
||||
raise RuntimeError("simulated encode failure")
|
||||
|
||||
vae = RaisingVAEModel()
|
||||
x = torch.zeros((1, 3, 12, 64, 64), dtype=torch.float32)
|
||||
|
||||
with pytest.raises(RuntimeError, match="simulated encode failure"):
|
||||
tiled_vae(
|
||||
x,
|
||||
vae,
|
||||
tile_size=(64, 64),
|
||||
tile_overlap=(0, 0),
|
||||
temporal_size=0,
|
||||
temporal_overlap=0,
|
||||
encode=True,
|
||||
)
|
||||
|
||||
assert vae.slicing_sample_min_size == 4
|
||||
|
||||
|
||||
def test_tiled_vae_encode_uses_tensor_return_without_indexing():
|
||||
class TensorEncodeVAEModel(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.slicing_sample_min_size = 4
|
||||
self.spatial_downsample_factor = 8
|
||||
self.temporal_downsample_factor = 4
|
||||
self.device = torch.device("cpu")
|
||||
self._dummy = torch.nn.Parameter(torch.zeros(1, dtype=torch.float32))
|
||||
self.calls = []
|
||||
|
||||
def encode(self, t_chunk):
|
||||
self.calls.append(tuple(t_chunk.shape))
|
||||
b, _, _, h, w = t_chunk.shape
|
||||
return torch.ones((b, _LATENT_CHANNELS, 1, h // 8, w // 8), dtype=t_chunk.dtype)
|
||||
|
||||
vae = TensorEncodeVAEModel()
|
||||
x = torch.zeros((2, 3, 1, 64, 64), dtype=torch.float32)
|
||||
|
||||
out = tiled_vae(
|
||||
x,
|
||||
vae,
|
||||
tile_size=(64, 64),
|
||||
tile_overlap=(0, 0),
|
||||
temporal_size=0,
|
||||
temporal_overlap=0,
|
||||
encode=True,
|
||||
)
|
||||
|
||||
assert vae.calls == [(2, 3, 1, 64, 64)]
|
||||
assert tuple(out.shape) == (2, _LATENT_CHANNELS, 1, 8, 8)
|
||||
|
||||
|
||||
def test_tiled_vae_preserves_compute_dtype_with_different_parameter_dtype():
|
||||
class DummyVAE(nn.Module):
|
||||
spatial_downsample_factor = 8
|
||||
temporal_downsample_factor = 4
|
||||
slicing_sample_min_size = 8
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.device = torch.device("cpu")
|
||||
self._dummy = nn.Parameter(torch.zeros(1, dtype=torch.float16))
|
||||
self.input_dtype = None
|
||||
|
||||
def encode(self, t_chunk):
|
||||
self.input_dtype = t_chunk.dtype
|
||||
b, _, _, h, w = t_chunk.shape
|
||||
return torch.ones((b, _LATENT_CHANNELS, 1, h // 8, w // 8), dtype=t_chunk.dtype)
|
||||
|
||||
vae = DummyVAE()
|
||||
x = torch.zeros((1, 3, 1, 64, 64), dtype=torch.float32)
|
||||
|
||||
tiled_vae(x, vae, tile_size=(64, 64), tile_overlap=(16, 16), encode=True)
|
||||
|
||||
assert vae.input_dtype == torch.float32
|
||||
|
||||
|
||||
def test_tiled_vae_preserves_input_dtype_on_single_tile():
|
||||
class FloatOutputVAEModel(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.slicing_sample_min_size = 4
|
||||
self.spatial_downsample_factor = 8
|
||||
self.temporal_downsample_factor = 4
|
||||
self.device = torch.device("cpu")
|
||||
self._dummy = torch.nn.Parameter(torch.zeros(1, dtype=torch.float32))
|
||||
|
||||
def encode(self, t_chunk):
|
||||
b, _, _, h, w = t_chunk.shape
|
||||
return torch.ones((b, _LATENT_CHANNELS, 1, h // 8, w // 8), dtype=torch.float32)
|
||||
|
||||
out = tiled_vae(
|
||||
torch.zeros((1, 3, 1, 64, 64), dtype=torch.float16),
|
||||
FloatOutputVAEModel(),
|
||||
tile_size=(64, 64),
|
||||
tile_overlap=(0, 0),
|
||||
temporal_size=0,
|
||||
temporal_overlap=0,
|
||||
encode=True,
|
||||
)
|
||||
|
||||
assert out.dtype == torch.float16
|
||||
|
||||
|
||||
class _SlicingDecodeVAE(nn.Module):
|
||||
def __init__(self, slicing_latent_min_size):
|
||||
super().__init__()
|
||||
self.slicing_latent_min_size = slicing_latent_min_size
|
||||
self.spatial_downsample_factor = 8
|
||||
self.temporal_downsample_factor = 4
|
||||
self.device = torch.device("cpu")
|
||||
self.use_slicing = True
|
||||
self._dummy = nn.Parameter(torch.zeros(1, dtype=torch.float32))
|
||||
self.decode_min_sizes = []
|
||||
self.memory_states = []
|
||||
|
||||
def decode_(self, z):
|
||||
self.decode_min_sizes.append(self.slicing_latent_min_size)
|
||||
return vae_mod.VideoAutoencoderKL.slicing_decode(self, z)
|
||||
|
||||
def _decode(self, z, memory_state=MemoryState.DISABLED, memory_cache=None):
|
||||
self.memory_states.append(memory_state)
|
||||
x = z[:, :1].repeat(
|
||||
1,
|
||||
3,
|
||||
1,
|
||||
self.spatial_downsample_factor,
|
||||
self.spatial_downsample_factor,
|
||||
)
|
||||
return x
|
||||
|
||||
|
||||
def test_decode_tiled_vae_maps_temporal_args_to_latent_slicing_min_size():
|
||||
vae = _SlicingDecodeVAE(slicing_latent_min_size=2)
|
||||
z = torch.arange(
|
||||
_LATENT_CHANNELS * 5 * 8 * 8,
|
||||
dtype=torch.float32,
|
||||
).reshape(1, _LATENT_CHANNELS, 5, 8, 8)
|
||||
|
||||
tiled_vae(
|
||||
z,
|
||||
vae,
|
||||
tile_size=(64, 64),
|
||||
tile_overlap=(0, 0),
|
||||
temporal_size=12,
|
||||
temporal_overlap=4,
|
||||
encode=False,
|
||||
)
|
||||
|
||||
assert vae.decode_min_sizes == [2]
|
||||
assert vae.memory_states == [MemoryState.INITIALIZING, MemoryState.ACTIVE]
|
||||
assert vae.slicing_latent_min_size == 2
|
||||
|
||||
wrapper = vae_mod.VideoAutoencoderKLWrapper.__new__(
|
||||
vae_mod.VideoAutoencoderKLWrapper
|
||||
)
|
||||
nn.Module.__init__(wrapper)
|
||||
seedvr2_tiling = {
|
||||
"enable_tiling": True,
|
||||
"tile_size": (64, 64),
|
||||
"tile_overlap": (0, 0),
|
||||
"temporal_size": 8,
|
||||
"temporal_overlap": 7,
|
||||
}
|
||||
|
||||
captured = {}
|
||||
|
||||
def _fake_tiled_vae(latent, model, **kwargs):
|
||||
captured.update(kwargs)
|
||||
return torch.zeros(1, 3, 1, 16, 16)
|
||||
|
||||
with patch.object(vae_mod, "tiled_vae", side_effect=_fake_tiled_vae):
|
||||
wrapper.decode(torch.zeros(1, _LATENT_CHANNELS, 2, 2), seedvr2_tiling=seedvr2_tiling)
|
||||
|
||||
assert captured["temporal_overlap"] == 7
|
||||
|
||||
|
||||
def _force_oom(*a, **k):
|
||||
raise torch.cuda.OutOfMemoryError("forced OOM for dispatcher test")
|
||||
|
||||
|
||||
def _make_vae(first_stage_model, latent_channels, latent_dim):
|
||||
vae = sd_mod.VAE.__new__(sd_mod.VAE)
|
||||
vae.first_stage_model = first_stage_model
|
||||
vae.patcher = MagicMock()
|
||||
vae.patcher.get_free_memory = MagicMock(return_value=8 * 1024 * 1024 * 1024)
|
||||
vae.device = vae.output_device = torch.device("cpu")
|
||||
vae.vae_dtype = torch.float32
|
||||
vae.disable_offload = True
|
||||
vae.extra_1d_channel = None
|
||||
vae.upscale_ratio = vae.downscale_ratio = 8
|
||||
vae.upscale_index_formula = vae.downscale_index_formula = None
|
||||
vae.output_channels = 3
|
||||
vae.latent_channels = latent_channels
|
||||
vae.latent_dim = latent_dim
|
||||
vae.vae_output_dtype = lambda: torch.float32
|
||||
vae.spacial_compression_decode = lambda: 8
|
||||
vae.handles_tiling = isinstance(first_stage_model, seedvr_vae_mod.VideoAutoencoderKLWrapper)
|
||||
vae.format_encoded = None
|
||||
vae.process_input = lambda x: x
|
||||
vae.process_output = lambda x: x
|
||||
vae.throw_exception_if_invalid = lambda: None
|
||||
vae.memory_used_decode = lambda *a, **k: 1
|
||||
return vae
|
||||
|
||||
|
||||
def _dispatch(vae, samples, seedvr2_call, generic_call, patch_wrapper_decode):
|
||||
mm = sd_mod.model_management
|
||||
with ExitStack() as stack:
|
||||
stack.enter_context(patch.object(mm, "raise_non_oom", lambda e: None))
|
||||
stack.enter_context(patch.object(mm, "load_models_gpu", lambda *a, **k: None))
|
||||
stack.enter_context(patch.object(mm, "soft_empty_cache", lambda: None))
|
||||
stack.enter_context(patch.object(sd_mod.VAE, "_decode_tiled_owned", seedvr2_call))
|
||||
stack.enter_context(patch.object(sd_mod.VAE, "decode_tiled_", generic_call))
|
||||
if patch_wrapper_decode:
|
||||
stack.enter_context(patch.object(
|
||||
seedvr_vae_mod.VideoAutoencoderKLWrapper, "decode",
|
||||
side_effect=_force_oom))
|
||||
vae.decode(samples)
|
||||
|
||||
|
||||
def test_4d_seedvr2_latent_routes_to_owned_decode_tiled():
|
||||
wrapper = seedvr_vae_mod.VideoAutoencoderKLWrapper.__new__(
|
||||
seedvr_vae_mod.VideoAutoencoderKLWrapper)
|
||||
vae = _make_vae(wrapper, latent_channels=_LATENT_CHANNELS, latent_dim=3)
|
||||
seedvr2_call = MagicMock(return_value=torch.zeros(1, 3, 9, 64, 64))
|
||||
generic_call = MagicMock(return_value=torch.zeros(1, 3, 64, 64))
|
||||
_dispatch(vae, torch.zeros(1, _LATENT_CHANNELS * 3, 8, 8), seedvr2_call, generic_call, True)
|
||||
assert seedvr2_call.call_count == 1
|
||||
assert generic_call.call_count == 0
|
||||
|
||||
|
||||
def test_4d_non_seedvr2_latent_still_routes_to_generic_decode_tiled():
|
||||
first_stage = MagicMock()
|
||||
first_stage.decode = MagicMock(side_effect=_force_oom)
|
||||
vae = _make_vae(first_stage, latent_channels=4, latent_dim=2)
|
||||
seedvr2_call = MagicMock(return_value=torch.zeros(1, 3, 9, 64, 64))
|
||||
generic_call = MagicMock(return_value=torch.zeros(1, 3, 64, 64))
|
||||
_dispatch(vae, torch.zeros(1, 4, 8, 8), seedvr2_call, generic_call, False)
|
||||
assert generic_call.call_count == 1
|
||||
assert seedvr2_call.call_count == 0
|
||||
|
||||
|
||||
def _populate_common_vae_attrs_fallback(vae):
|
||||
vae.patcher = MagicMock()
|
||||
vae.patcher.get_free_memory = MagicMock(return_value=8 * 1024 * 1024 * 1024)
|
||||
vae.device = torch.device("cpu")
|
||||
vae.output_device = torch.device("cpu")
|
||||
vae.vae_dtype = torch.float32
|
||||
vae.disable_offload = True
|
||||
vae.extra_1d_channel = None
|
||||
vae.upscale_ratio = 8
|
||||
vae.upscale_index_formula = None
|
||||
vae.output_channels = 3
|
||||
vae.latent_channels = _LATENT_CHANNELS
|
||||
vae.latent_dim = 3
|
||||
vae.downscale_ratio = 8
|
||||
vae.downscale_index_formula = None
|
||||
vae.not_video = False
|
||||
vae.crop_input = False
|
||||
vae.pad_channel_value = None
|
||||
vae.handles_tiling = isinstance(vae.first_stage_model, seedvr_vae_mod.VideoAutoencoderKLWrapper)
|
||||
vae.format_encoded = None
|
||||
|
||||
vae.vae_output_dtype = lambda: torch.float32
|
||||
vae.spacial_compression_encode = lambda: 8
|
||||
vae.process_input = lambda x: x
|
||||
vae.process_output = lambda x: x
|
||||
vae.throw_exception_if_invalid = lambda: None
|
||||
vae.memory_used_encode = lambda *a, **k: 1
|
||||
|
||||
|
||||
def _make_seedvr2_vae_fallback():
|
||||
vae = sd_mod.VAE.__new__(sd_mod.VAE)
|
||||
wrapper = seedvr_vae_mod.VideoAutoencoderKLWrapper.__new__(
|
||||
seedvr_vae_mod.VideoAutoencoderKLWrapper
|
||||
)
|
||||
vae.first_stage_model = wrapper
|
||||
_populate_common_vae_attrs_fallback(vae)
|
||||
return vae
|
||||
|
||||
|
||||
def _make_non_seedvr2_vae_fallback():
|
||||
vae = sd_mod.VAE.__new__(sd_mod.VAE)
|
||||
vae.first_stage_model = MagicMock()
|
||||
_populate_common_vae_attrs_fallback(vae)
|
||||
return vae
|
||||
|
||||
|
||||
def _force_regular_encode_oom(*args, **kwargs):
|
||||
raise torch.cuda.OutOfMemoryError("forced OOM for dispatcher test")
|
||||
|
||||
|
||||
def test_seedvr2_3d_routes_to_owned_encode_tiled_on_oom():
|
||||
vae = _make_seedvr2_vae_fallback()
|
||||
pixel_samples = torch.zeros((1, 8, 64, 64, 3))
|
||||
|
||||
seedvr2_call = MagicMock(return_value=torch.zeros(1, _LATENT_CHANNELS, 2, 8, 8))
|
||||
generic_call = MagicMock(return_value=torch.zeros(1, _LATENT_CHANNELS, 2, 8, 8))
|
||||
|
||||
with patch.object(sd_mod.model_management, "raise_non_oom",
|
||||
lambda e: None), \
|
||||
patch.object(sd_mod.model_management, "load_models_gpu",
|
||||
lambda *a, **k: None), \
|
||||
patch.object(sd_mod.model_management, "soft_empty_cache",
|
||||
lambda: None), \
|
||||
patch.object(seedvr_vae_mod.VideoAutoencoderKLWrapper, "encode",
|
||||
side_effect=_force_regular_encode_oom), \
|
||||
patch.object(sd_mod.VAE, "_encode_tiled_owned", seedvr2_call), \
|
||||
patch.object(sd_mod.VAE, "encode_tiled_3d", generic_call):
|
||||
vae.encode(pixel_samples)
|
||||
|
||||
assert seedvr2_call.call_count == 1, (
|
||||
f"Expected _encode_tiled_owned to be called once for a SeedVR2 3D "
|
||||
f"input under OOM fallback; got {seedvr2_call.call_count} calls."
|
||||
)
|
||||
assert generic_call.call_count == 0, (
|
||||
f"encode_tiled_3d must NOT be called for a SeedVR2 input; got "
|
||||
f"{generic_call.call_count} calls."
|
||||
)
|
||||
|
||||
|
||||
def test_non_seedvr2_encode_tiled_3d_default_overlap_is_concrete():
|
||||
vae = _make_non_seedvr2_vae_fallback()
|
||||
vae.downscale_ratio = (lambda a: max(1, a // 4), 8, 8)
|
||||
vae.upscale_ratio = (lambda a: a * 4, 8, 8)
|
||||
generic_call = MagicMock(return_value=torch.zeros(1, _LATENT_CHANNELS, 2, 8, 8))
|
||||
pixel_samples = torch.zeros((1, 8, 64, 64, 3))
|
||||
|
||||
with patch.object(sd_mod.model_management, "load_models_gpu",
|
||||
lambda *a, **k: None), \
|
||||
patch.object(sd_mod.VAE, "encode_tiled_3d", generic_call):
|
||||
vae.encode_tiled(pixel_samples)
|
||||
|
||||
assert generic_call.call_args.kwargs["overlap"] == (1, 64, 64)
|
||||
@@ -11,11 +11,6 @@ from comfy_api.feature_flags import (
|
||||
_coerce_flag_value,
|
||||
_parse_cli_feature_flags,
|
||||
)
|
||||
from comfy.comfy_api_env import (
|
||||
environment_overrides_for_base,
|
||||
get_environment_overrides,
|
||||
normalize_comfy_api_base,
|
||||
)
|
||||
|
||||
|
||||
class TestFeatureFlags:
|
||||
@@ -188,65 +183,3 @@ class TestCliFeatureFlagRegistry:
|
||||
assert "type" in info, f"{key} missing 'type'"
|
||||
assert "default" in info, f"{key} missing 'default'"
|
||||
assert "description" in info, f"{key} missing 'description'"
|
||||
|
||||
|
||||
class TestComfyApiEnv:
|
||||
"""--comfy-api-base staging-tier detection + testenv main-host -> -registry rewrite."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"url, expected",
|
||||
[
|
||||
# testenv friendly main host -> comfy-api -registry sibling (slash trimmed)
|
||||
("https://pr-4398.testenvs.comfy.org", "https://pr-4398-registry.testenvs.comfy.org"),
|
||||
("https://pr-4398.testenvs.comfy.org/", "https://pr-4398-registry.testenvs.comfy.org"),
|
||||
("https://pr-4398-registry.testenvs.comfy.org", "https://pr-4398-registry.testenvs.comfy.org"),
|
||||
# staging + everything else -> unchanged (no -registry split)
|
||||
("https://stagingapi.comfy.org", "https://stagingapi.comfy.org"),
|
||||
("https://api.comfy.org", "https://api.comfy.org"),
|
||||
("https://pr-1.testenvs.comfy.org.evil.com", "https://pr-1.testenvs.comfy.org.evil.com"),
|
||||
("", ""),
|
||||
],
|
||||
)
|
||||
def test_normalize_comfy_api_base(self, url, expected):
|
||||
assert normalize_comfy_api_base(url) == expected
|
||||
|
||||
def test_config_for_staging_tier_else_none(self):
|
||||
# ephemeral testenv: friendly main host -> -registry, staging platform, dev Firebase env
|
||||
eph = environment_overrides_for_base("https://pr-1234.testenvs.comfy.org/")
|
||||
assert eph["comfy_api_base_url"] == "https://pr-1234-registry.testenvs.comfy.org"
|
||||
assert eph["comfy_platform_base_url"] == "https://stagingplatform.comfy.org"
|
||||
assert eph["firebase_env"] == "dev"
|
||||
# staging api host: emitted as-is
|
||||
stg = environment_overrides_for_base("https://stagingapi.comfy.org")
|
||||
assert stg["comfy_api_base_url"] == "https://stagingapi.comfy.org"
|
||||
assert stg["comfy_platform_base_url"] == "https://stagingplatform.comfy.org"
|
||||
assert stg["firebase_env"] == "dev"
|
||||
# prod / unknown: nothing
|
||||
assert environment_overrides_for_base("https://api.comfy.org") is None
|
||||
|
||||
def test_environment_overrides_only_for_staging_tier(self, monkeypatch):
|
||||
def set_base(url):
|
||||
monkeypatch.setattr(
|
||||
"comfy.comfy_api_env.args",
|
||||
type("Args", (), {"comfy_api_base": url})(),
|
||||
)
|
||||
|
||||
# The overrides merged into the HTTP /features response are present for staging-tier bases...
|
||||
set_base("https://stagingapi.comfy.org")
|
||||
assert "comfy_api_base_url" in get_environment_overrides()
|
||||
set_base("https://pr-7.testenvs.comfy.org")
|
||||
assert "comfy_api_base_url" in get_environment_overrides()
|
||||
# ...but never for prod.
|
||||
set_base("https://api.comfy.org")
|
||||
assert get_environment_overrides() is None
|
||||
|
||||
def test_server_features_never_carry_env_overrides(self, monkeypatch):
|
||||
"""The WebSocket capability handshake must stay free of routing keys."""
|
||||
monkeypatch.setattr(
|
||||
"comfy.comfy_api_env.args",
|
||||
type("Args", (), {"comfy_api_base": "https://pr-7.testenvs.comfy.org"})(),
|
||||
)
|
||||
features = get_server_features()
|
||||
assert "comfy_api_base_url" not in features
|
||||
assert "comfy_platform_base_url" not in features
|
||||
assert "firebase_env" not in features
|
||||
|
||||
@@ -281,41 +281,6 @@ class TestAsyncNodes:
|
||||
# Verify the sync error was caught even though async was running
|
||||
assert 'prompt_id' in e.args[0]
|
||||
|
||||
def test_async_sibling_completes_after_error(self, client: ComfyClient, builder: GraphBuilder):
|
||||
g = builder
|
||||
image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
|
||||
error_node = g.node("TestAsyncError", value=image.out(0), error_after=0.05)
|
||||
sleep_node = g.node("TestSleep", value=image.out(0), seconds=0.1)
|
||||
g.node("PreviewImage", images=error_node.out(0))
|
||||
successful_output = g.node("SaveImage", images=sleep_node.out(0))
|
||||
|
||||
result = client.run(g, node_failure_policy="continue_independent")
|
||||
|
||||
assert result.did_run(error_node)
|
||||
assert result.did_run(sleep_node)
|
||||
assert result.was_executed(successful_output)
|
||||
assert len(result.get_images(successful_output)) == 1
|
||||
assert result.execution_success['completion_status'] == 'partial_success'
|
||||
assert result.node_errors[0]['node_id'] == error_node.id
|
||||
|
||||
def test_async_sibling_completes_after_multiple_errors(self, client: ComfyClient, builder: GraphBuilder):
|
||||
g = builder
|
||||
image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
|
||||
error1 = g.node("TestAsyncError", value=image.out(0), error_after=0.02)
|
||||
error2 = g.node("TestAsyncError", value=image.out(0), error_after=0.04)
|
||||
sleep_node = g.node("TestSleep", value=image.out(0), seconds=0.06)
|
||||
g.node("PreviewImage", images=error1.out(0))
|
||||
g.node("PreviewImage", images=error2.out(0))
|
||||
successful_output = g.node("SaveImage", images=sleep_node.out(0))
|
||||
|
||||
result = client.run(g, node_failure_policy="continue_independent")
|
||||
|
||||
assert {error['node_id'] for error in result.node_errors} == {error1.id, error2.id}
|
||||
assert result.did_run(sleep_node)
|
||||
assert result.was_executed(successful_output)
|
||||
assert result.execution_success['completion_status'] == 'partial_success'
|
||||
assert result.execution_success['execution_error_count'] == 2
|
||||
|
||||
# Edge Cases
|
||||
|
||||
def test_async_with_execution_blocker(self, client: ComfyClient, builder: GraphBuilder):
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
from collections import Counter
|
||||
from io import BytesIO
|
||||
import numpy
|
||||
from PIL import Image
|
||||
@@ -14,35 +13,8 @@ import uuid
|
||||
import urllib.request
|
||||
import urllib.parse
|
||||
import urllib.error
|
||||
from comfy_execution.graph import DynamicPrompt, ExecutionList
|
||||
from comfy_execution.graph_utils import GraphBuilder, Node
|
||||
|
||||
|
||||
def test_execution_list_transient_cache_is_prompt_scoped():
|
||||
class OutputCache:
|
||||
def __init__(self):
|
||||
self.values = {}
|
||||
self.set_calls = []
|
||||
|
||||
def get_local(self, node_id):
|
||||
return self.values.get(node_id)
|
||||
|
||||
def set_local(self, node_id, value):
|
||||
self.set_calls.append((node_id, value))
|
||||
self.values[node_id] = value
|
||||
|
||||
output_cache = OutputCache()
|
||||
execution_list = ExecutionList(DynamicPrompt({}), output_cache)
|
||||
failure_entry = object()
|
||||
|
||||
execution_list.cache_update("failed", failure_entry, transient=True)
|
||||
execution_list.cache_link("failed", "late_consumer")
|
||||
|
||||
assert execution_list.is_cached("failed")
|
||||
assert execution_list.get_cache("failed", "late_consumer") is failure_entry
|
||||
assert output_cache.set_calls == []
|
||||
assert not ExecutionList(DynamicPrompt({}), output_cache).is_cached("failed")
|
||||
|
||||
def run_warmup(client, prefix="warmup"):
|
||||
"""Run a simple workflow to warm up the server."""
|
||||
warmup_g = GraphBuilder(prefix=prefix)
|
||||
@@ -56,9 +28,6 @@ class RunResult:
|
||||
self.runs: Dict[str,bool] = {}
|
||||
self.cached: Dict[str,bool] = {}
|
||||
self.prompt_id: str = prompt_id
|
||||
self.node_errors = []
|
||||
self.execution_success = None
|
||||
self.run_counts: Dict[str, int] = {}
|
||||
|
||||
def get_output(self, node: Node):
|
||||
return self.outputs.get(node.id, None)
|
||||
@@ -97,12 +66,10 @@ class ComfyClient:
|
||||
ws.connect("ws://{}/ws?clientId={}".format(self.server_address, self.client_id))
|
||||
self.ws = ws
|
||||
|
||||
def queue_prompt(self, prompt, partial_execution_targets=None, node_failure_policy=None):
|
||||
def queue_prompt(self, prompt, partial_execution_targets=None):
|
||||
p = {"prompt": prompt, "client_id": self.client_id}
|
||||
if partial_execution_targets is not None:
|
||||
p["partial_execution_targets"] = partial_execution_targets
|
||||
if node_failure_policy is not None:
|
||||
p["node_failure_policy"] = node_failure_policy
|
||||
data = json.dumps(p).encode('utf-8')
|
||||
req = urllib.request.Request("http://{}/prompt".format(self.server_address), data=data)
|
||||
return json.loads(urllib.request.urlopen(req).read())
|
||||
@@ -166,13 +133,13 @@ class ComfyClient:
|
||||
def set_test_name(self, name):
|
||||
self.test_name = name
|
||||
|
||||
def run(self, graph, partial_execution_targets=None, node_failure_policy=None):
|
||||
def run(self, graph, partial_execution_targets=None):
|
||||
prompt = graph.finalize()
|
||||
for node in graph.nodes.values():
|
||||
if node.class_type == 'SaveImage':
|
||||
node.inputs['filename_prefix'] = self.test_name
|
||||
|
||||
prompt_id = self.queue_prompt(prompt, partial_execution_targets, node_failure_policy)['prompt_id']
|
||||
prompt_id = self.queue_prompt(prompt, partial_execution_targets)['prompt_id']
|
||||
result = RunResult(prompt_id)
|
||||
while True:
|
||||
out = self.ws.recv()
|
||||
@@ -185,15 +152,8 @@ class ComfyClient:
|
||||
if data['node'] is None:
|
||||
break
|
||||
result.runs[data['node']] = True
|
||||
result.run_counts[data['node']] = result.run_counts.get(data['node'], 0) + 1
|
||||
elif message['type'] == 'execution_error':
|
||||
raise Exception(message['data'])
|
||||
elif message['type'] == 'execution_node_error':
|
||||
if message['data']['prompt_id'] == prompt_id:
|
||||
result.node_errors.append(message['data'])
|
||||
elif message['type'] == 'execution_success':
|
||||
if message['data']['prompt_id'] == prompt_id:
|
||||
result.execution_success = message['data']
|
||||
elif message['type'] == 'execution_cached':
|
||||
if message['data']['prompt_id'] == prompt_id:
|
||||
cached_nodes = message['data'].get('nodes', [])
|
||||
@@ -345,299 +305,6 @@ class TestExecution:
|
||||
except Exception as e:
|
||||
assert 'prompt_id' in e.args[0], f"Did not get back a proper error message: {e}"
|
||||
|
||||
def test_continue_independent_after_error(self, client: ComfyClient, builder: GraphBuilder):
|
||||
g = builder
|
||||
image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
|
||||
error_node = g.node("TestSyncError", value=image.out(0))
|
||||
blocked_output = g.node("PreviewImage", images=error_node.out(0))
|
||||
successful_output = g.node("SaveImage", images=image.out(0))
|
||||
|
||||
result = client.run(g, node_failure_policy="continue_independent")
|
||||
|
||||
assert result.did_run(error_node)
|
||||
assert result.was_executed(successful_output)
|
||||
assert len(result.get_images(successful_output)) == 1
|
||||
assert len(result.node_errors) == 1
|
||||
assert result.node_errors[0]['node_id'] == error_node.id
|
||||
assert result.execution_success['completion_status'] == 'partial_success'
|
||||
assert result.execution_success['has_errors'] is True
|
||||
assert result.execution_success['execution_error_count'] == 1
|
||||
assert blocked_output.id in result.execution_success['blocked_output_node_ids']
|
||||
assert successful_output.id in result.execution_success['successful_output_node_ids']
|
||||
history = client.get_history(result.prompt_id)[result.prompt_id]
|
||||
assert '_node_failure_policy' not in history['prompt'][3]
|
||||
|
||||
retry = client.run(g, node_failure_policy="continue_independent")
|
||||
assert retry.did_run(error_node), "Failed nodes must be retried on a new prompt"
|
||||
assert len(retry.node_errors) == 1
|
||||
|
||||
def test_continue_independent_reuses_failed_node_for_late_lazy_link(self, client: ComfyClient, builder: GraphBuilder):
|
||||
g = builder
|
||||
image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
|
||||
error_node = g.node("TestSyncError", value=image.out(0))
|
||||
g.node("PreviewImage", images=error_node.out(0))
|
||||
|
||||
mask = g.node("StubMask", value=0.0, height=32, width=32, batch_size=1)
|
||||
unused_image = g.node("StubImage", content="WHITE", height=32, width=32, batch_size=1)
|
||||
lazy_mix = g.node("TestLazyMixImages", image1=error_node.out(0), image2=unused_image.out(0), mask=mask.out(0))
|
||||
g.node("PreviewImage", images=lazy_mix.out(0))
|
||||
successful_output = g.node("SaveImage", images=image.out(0))
|
||||
|
||||
result = client.run(g, node_failure_policy="continue_independent")
|
||||
|
||||
assert result.run_counts[error_node.id] == 1
|
||||
assert lazy_mix.id in result.execution_success['blocked_node_ids']
|
||||
assert successful_output.id in result.execution_success['successful_output_node_ids']
|
||||
assert result.execution_success['completion_status'] == 'partial_success'
|
||||
|
||||
def test_continue_independent_handles_mixed_dynamic_results(self, client: ComfyClient, builder: GraphBuilder):
|
||||
g = builder
|
||||
zero = g.node("StubInt", value=0)
|
||||
one = g.node("StubInt", value=1)
|
||||
values = g.node("TestMakeListNode", value1=zero.out(0), value2=one.out(0))
|
||||
mixed = g.node("TestMixedExpansionFailure", value=values.out(0))
|
||||
output = g.node("PreviewImage", images=mixed.out(0))
|
||||
|
||||
result = client.run(g, node_failure_policy="continue_independent")
|
||||
|
||||
assert len(result.node_errors) == 1
|
||||
assert result.node_errors[0]['node_id'] == mixed.id
|
||||
assert len(result.get_images(output)) == 1
|
||||
assert output.id in result.execution_success['successful_output_node_ids']
|
||||
assert output.id not in result.execution_success['blocked_output_node_ids']
|
||||
assert result.execution_success['completion_status'] == 'partial_success'
|
||||
|
||||
retry = client.run(g, node_failure_policy="continue_independent")
|
||||
assert retry.did_run(mixed), "Failure-tainted dynamic parents must not be reused from cache"
|
||||
assert len(retry.node_errors) == 1
|
||||
|
||||
def test_continue_independent_keeps_oom_terminal(self, client: ComfyClient, builder: GraphBuilder):
|
||||
g = builder
|
||||
image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
|
||||
oom_node = g.node("TestOOMError", value=image.out(0))
|
||||
g.node("PreviewImage", images=oom_node.out(0))
|
||||
g.node("SaveImage", images=image.out(0))
|
||||
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
client.run(g, node_failure_policy="continue_independent")
|
||||
|
||||
assert exc_info.value.args[0]['node_id'] == oom_node.id
|
||||
assert exc_info.value.args[0]['exception_type'] == 'torch.OutOfMemoryError'
|
||||
|
||||
def test_continue_independent_accepts_cached_output(self, client: ComfyClient, builder: GraphBuilder, server):
|
||||
g = builder
|
||||
image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
|
||||
error_node = g.node("TestSyncError", value=image.out(0))
|
||||
g.node("PreviewImage", images=error_node.out(0))
|
||||
successful_output = g.node("SaveImage", images=image.out(0))
|
||||
|
||||
client.run(g, partial_execution_targets=[successful_output.id])
|
||||
result = client.run(g, node_failure_policy="continue_independent")
|
||||
|
||||
if server["should_cache_results"]:
|
||||
assert result.was_cached(successful_output)
|
||||
assert len(result.node_errors) == 1
|
||||
assert result.execution_success['completion_status'] == 'partial_success'
|
||||
assert successful_output.id in result.execution_success['successful_output_node_ids']
|
||||
|
||||
def test_continue_independent_keeps_executor_errors_terminal(self, client: ComfyClient, builder: GraphBuilder):
|
||||
g = builder
|
||||
image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
|
||||
bad = g.node("TestMalformedExpansion", value=image.out(0))
|
||||
g.node("PreviewImage", images=bad.out(0))
|
||||
g.node("SaveImage", images=image.out(0))
|
||||
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
client.run(g, node_failure_policy="continue_independent")
|
||||
|
||||
assert exc_info.value.args[0]['node_id'] == bad.id
|
||||
assert exc_info.value.args[0]['exception_type'] == 'KeyError'
|
||||
|
||||
def test_continue_independent_malformed_result_terminal(self, client: ComfyClient, builder: GraphBuilder):
|
||||
g = builder
|
||||
image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
|
||||
bad = g.node("TestMalformedResult", value=image.out(0))
|
||||
g.node("PreviewImage", images=bad.out(0))
|
||||
g.node("SaveImage", images=image.out(0))
|
||||
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
client.run(g, node_failure_policy="continue_independent")
|
||||
|
||||
assert exc_info.value.args[0]['node_id'] == bad.id
|
||||
assert exc_info.value.args[0]['exception_type'] == 'TypeError'
|
||||
|
||||
def test_continue_independent_cyclic_expansion_reports_cycle(self, client: ComfyClient, builder: GraphBuilder):
|
||||
g = builder
|
||||
image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
|
||||
cyclic = g.node("TestCyclicExpansion", value=image.out(0))
|
||||
g.node("PreviewImage", images=cyclic.out(0))
|
||||
g.node("SaveImage", images=image.out(0))
|
||||
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
client.run(g, node_failure_policy="continue_independent")
|
||||
|
||||
assert exc_info.value.args[0]['exception_type'] == 'graph.DependencyCycleError'
|
||||
|
||||
def test_continue_independent_async_output_partial_failure(self, client: ComfyClient, builder: GraphBuilder):
|
||||
g = builder
|
||||
zero = g.node("StubInt", value=0)
|
||||
one = g.node("StubInt", value=1)
|
||||
values = g.node("TestMakeListNode", value1=zero.out(0), value2=one.out(0))
|
||||
mixed = g.node("TestMixedExpansionFailure", value=values.out(0))
|
||||
async_output = g.node("TestAsyncOutput", value=mixed.out(0), seconds=0.1)
|
||||
|
||||
result = client.run(g, node_failure_policy="continue_independent")
|
||||
|
||||
assert len(result.node_errors) == 1
|
||||
assert result.execution_success['completion_status'] == 'partial_success'
|
||||
assert async_output.id in result.execution_success['successful_output_node_ids']
|
||||
assert async_output.id not in result.execution_success['blocked_node_ids']
|
||||
|
||||
retry = client.run(g, node_failure_policy="continue_independent")
|
||||
assert retry.did_run(async_output), "Async outputs with failure-blocked invocations must be retried"
|
||||
|
||||
def test_continue_independent_failure_blocker_beats_user_blocker(self, client: ComfyClient, builder: GraphBuilder):
|
||||
g = builder
|
||||
image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
|
||||
error_node = g.node("TestSyncError", value=image.out(0))
|
||||
user_blocker = g.node("TestExecutionBlocker", input=image.out(0), block=True, verbose=False)
|
||||
combo = g.node("TestMakeListNode", value1=user_blocker.out(0), value2=error_node.out(0))
|
||||
g.node("PreviewImage", images=combo.out(0))
|
||||
successful_output = g.node("SaveImage", images=image.out(0))
|
||||
|
||||
result = client.run(g, node_failure_policy="continue_independent")
|
||||
|
||||
assert result.execution_success['completion_status'] == 'partial_success'
|
||||
assert combo.id in result.execution_success['blocked_node_ids']
|
||||
assert successful_output.id in result.execution_success['successful_output_node_ids']
|
||||
|
||||
retry = client.run(g, node_failure_policy="continue_independent")
|
||||
assert retry.did_run(combo), "Nodes blocked by a failure must be retried even when also user-blocked"
|
||||
|
||||
def test_continue_independent_retries_failed_expansion_side_branch(self, client: ComfyClient, builder: GraphBuilder):
|
||||
g = builder
|
||||
image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
|
||||
expander = g.node("TestExpansionWithFailingOutput", image=image.out(0))
|
||||
output = g.node("PreviewImage", images=expander.out(0))
|
||||
|
||||
result = client.run(g, node_failure_policy="continue_independent")
|
||||
|
||||
assert len(result.node_errors) == 1
|
||||
assert result.node_errors[0]['node_id'] == expander.id
|
||||
assert result.execution_success['completion_status'] == 'partial_success'
|
||||
assert len(result.get_images(output)) == 1
|
||||
|
||||
retry = client.run(g, node_failure_policy="continue_independent")
|
||||
assert retry.did_run(expander), "Expansion parents with failed side branches must be retried"
|
||||
assert len(retry.node_errors) == 1
|
||||
|
||||
def test_continue_independent_with_partial_targets(self, client: ComfyClient, builder: GraphBuilder):
|
||||
g = builder
|
||||
image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
|
||||
error_node = g.node("TestSyncError", value=image.out(0))
|
||||
failing_output = g.node("PreviewImage", images=error_node.out(0))
|
||||
successful_output = g.node("SaveImage", images=image.out(0))
|
||||
unselected_output = g.node("SaveImage", images=image.out(0))
|
||||
|
||||
result = client.run(
|
||||
g,
|
||||
partial_execution_targets=[failing_output.id, successful_output.id],
|
||||
node_failure_policy="continue_independent",
|
||||
)
|
||||
|
||||
assert len(result.node_errors) == 1
|
||||
assert result.execution_success['completion_status'] == 'partial_success'
|
||||
assert successful_output.id in result.execution_success['successful_output_node_ids']
|
||||
assert failing_output.id in result.execution_success['blocked_output_node_ids']
|
||||
assert not result.was_executed(unselected_output), "Unselected outputs must not execute"
|
||||
|
||||
def test_explicit_fail_fast_policy_matches_default(self, client: ComfyClient, builder: GraphBuilder):
|
||||
g = builder
|
||||
image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
|
||||
error_node = g.node("TestSyncError", value=image.out(0))
|
||||
g.node("PreviewImage", images=error_node.out(0))
|
||||
g.node("SaveImage", images=image.out(0))
|
||||
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
client.run(g, node_failure_policy="fail_fast")
|
||||
|
||||
assert exc_info.value.args[0]['node_id'] == error_node.id
|
||||
history = client.get_history(exc_info.value.args[0]['prompt_id'])
|
||||
entry = next(iter(history.values()))
|
||||
assert entry['status']['status_str'] == 'error'
|
||||
assert 'execution_summary' not in entry['status']
|
||||
|
||||
def test_continue_independent_failure_after_sibling_output(self, client: ComfyClient, builder: GraphBuilder):
|
||||
g = builder
|
||||
image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
|
||||
fast_output = g.node("TestAsyncOutput", value=image.out(0), seconds=0.05)
|
||||
slow_error = g.node("TestAsyncError", value=image.out(0), error_after=0.5)
|
||||
g.node("PreviewImage", images=slow_error.out(0))
|
||||
|
||||
result = client.run(g, node_failure_policy="continue_independent")
|
||||
|
||||
assert len(result.node_errors) == 1
|
||||
assert result.node_errors[0]['node_id'] == slow_error.id
|
||||
assert result.execution_success['completion_status'] == 'partial_success'
|
||||
assert fast_output.id in result.execution_success['successful_output_node_ids']
|
||||
|
||||
def test_continue_independent_never_stores_failed_outputs_externally(self, client: ComfyClient, builder: GraphBuilder, server):
|
||||
def record_stored_counts(prefix):
|
||||
record_graph = GraphBuilder(prefix=prefix)
|
||||
record = record_graph.node("TestCacheProviderRecord")
|
||||
record_result = client.run(record_graph)
|
||||
return Counter(record_result.get_output(record)['stored_class_types'])
|
||||
|
||||
baseline = record_stored_counts("cache_baseline")
|
||||
|
||||
g = builder
|
||||
# Unique 31x31 signature so StubImage executes fresh instead of hitting the cache
|
||||
image = g.node("StubImage", content="BLACK", height=31, width=31, batch_size=1)
|
||||
error_node = g.node("TestSyncError", value=image.out(0))
|
||||
g.node("PreviewImage", images=error_node.out(0))
|
||||
g.node("SaveImage", images=image.out(0))
|
||||
|
||||
result = client.run(g, node_failure_policy="continue_independent")
|
||||
assert result.execution_success['completion_status'] == 'partial_success'
|
||||
|
||||
delta = record_stored_counts("cache_record") - baseline
|
||||
if server["should_cache_results"]:
|
||||
assert delta["StubImage"] >= 1, "Successful outputs should reach external cache providers"
|
||||
assert delta["TestSyncError"] == 0, "Failed node outputs must never reach external cache providers"
|
||||
assert delta["PreviewImage"] == 0, "Failure-blocked outputs must never reach external cache providers"
|
||||
|
||||
def test_history_status_omits_summary_by_default(self, client: ComfyClient, builder: GraphBuilder):
|
||||
g = builder
|
||||
image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
|
||||
g.node("PreviewImage", images=image.out(0))
|
||||
|
||||
result = client.run(g)
|
||||
|
||||
history = client.get_history(result.prompt_id)[result.prompt_id]
|
||||
assert history['status']['status_str'] == 'success'
|
||||
assert 'execution_summary' not in history['status']
|
||||
|
||||
def test_continue_independent_fails_when_no_output_survives(self, client: ComfyClient, builder: GraphBuilder):
|
||||
g = builder
|
||||
image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
|
||||
error_node = g.node("TestSyncError", value=image.out(0))
|
||||
g.node("PreviewImage", images=error_node.out(0))
|
||||
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
client.run(g, node_failure_policy="continue_independent")
|
||||
|
||||
assert exc_info.value.args[0]['node_id'] == error_node.id
|
||||
|
||||
def test_invalid_node_failure_policy(self, client: ComfyClient, builder: GraphBuilder):
|
||||
g = builder
|
||||
image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
|
||||
g.node("PreviewImage", images=image.out(0))
|
||||
|
||||
with pytest.raises(urllib.error.HTTPError) as exc_info:
|
||||
client.queue_prompt(g.finalize(), node_failure_policy="continue_everything")
|
||||
|
||||
assert exc_info.value.code == 400
|
||||
|
||||
@pytest.mark.parametrize("test_value, expect_error", [
|
||||
(5, True),
|
||||
("foo", True),
|
||||
|
||||
@@ -500,76 +500,6 @@ class TestNormalizeHistoryItem:
|
||||
'extra_data': {'create_time': 1234567890, 'client_id': 'abc'},
|
||||
}
|
||||
|
||||
def test_missing_status(self):
|
||||
history_item = {
|
||||
'prompt': (
|
||||
5,
|
||||
'prompt-without-status',
|
||||
{'nodes': {}},
|
||||
{'create_time': 100},
|
||||
[],
|
||||
),
|
||||
'status': None,
|
||||
'outputs': {},
|
||||
}
|
||||
|
||||
job = normalize_history_item('prompt-without-status', history_item)
|
||||
|
||||
assert job['status'] == 'completed'
|
||||
assert 'completion_status' not in job
|
||||
|
||||
def test_partial_success_metadata_and_errors(self):
|
||||
node_error = {
|
||||
'prompt_id': 'prompt-partial',
|
||||
'node_id': '2',
|
||||
'node_type': 'TestSyncError',
|
||||
'exception_message': 'failed',
|
||||
'exception_type': 'RuntimeError',
|
||||
'traceback': [],
|
||||
'current_inputs': {},
|
||||
'current_outputs': [],
|
||||
'timestamp': 200,
|
||||
}
|
||||
history_item = {
|
||||
'prompt': (
|
||||
5,
|
||||
'prompt-partial',
|
||||
{'nodes': {}},
|
||||
{'create_time': 100},
|
||||
['3', '4'],
|
||||
),
|
||||
'status': {
|
||||
'status_str': 'success',
|
||||
'completed': True,
|
||||
'execution_summary': {
|
||||
'completion_status': 'partial_success',
|
||||
'has_errors': True,
|
||||
'execution_error_count': 1,
|
||||
},
|
||||
'messages': [
|
||||
('execution_start', {'prompt_id': 'prompt-partial', 'timestamp': 150}),
|
||||
('execution_node_error', node_error),
|
||||
('execution_success', {
|
||||
'prompt_id': 'prompt-partial',
|
||||
'completion_status': 'partial_success',
|
||||
'has_errors': True,
|
||||
'execution_error_count': 1,
|
||||
'timestamp': 300,
|
||||
}),
|
||||
],
|
||||
},
|
||||
'outputs': {'4': {'images': [{'filename': 'survived.png'}]}},
|
||||
}
|
||||
|
||||
job = normalize_history_item('prompt-partial', history_item, include_outputs=True)
|
||||
|
||||
assert job['status'] == 'completed'
|
||||
assert job['completion_status'] == 'partial_success'
|
||||
assert job['has_errors'] is True
|
||||
assert job['execution_error_count'] == 1
|
||||
assert job['execution_errors'] == [node_error]
|
||||
assert job['outputs']['4']['images'] == [{'filename': 'survived.png'}]
|
||||
|
||||
def test_include_outputs_normalizes_3d_strings(self):
|
||||
"""Detail view should transform string 3D filenames into file output dicts."""
|
||||
history_item = {
|
||||
|
||||
@@ -5,7 +5,6 @@ from .conditions import CONDITION_NODE_CLASS_MAPPINGS, CONDITION_NODE_DISPLAY_NA
|
||||
from .stubs import TEST_STUB_NODE_CLASS_MAPPINGS, TEST_STUB_NODE_DISPLAY_NAME_MAPPINGS
|
||||
from .async_test_nodes import ASYNC_TEST_NODE_CLASS_MAPPINGS, ASYNC_TEST_NODE_DISPLAY_NAME_MAPPINGS
|
||||
from .api_test_nodes import API_TEST_NODE_CLASS_MAPPINGS, API_TEST_NODE_DISPLAY_NAME_MAPPINGS
|
||||
from .cache_provider_test_nodes import CACHE_PROVIDER_TEST_NODE_CLASS_MAPPINGS, CACHE_PROVIDER_TEST_NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
# NODE_CLASS_MAPPINGS = GENERAL_NODE_CLASS_MAPPINGS.update(COMPONENT_NODE_CLASS_MAPPINGS)
|
||||
# NODE_DISPLAY_NAME_MAPPINGS = GENERAL_NODE_DISPLAY_NAME_MAPPINGS.update(COMPONENT_NODE_DISPLAY_NAME_MAPPINGS)
|
||||
@@ -18,7 +17,6 @@ NODE_CLASS_MAPPINGS.update(CONDITION_NODE_CLASS_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(TEST_STUB_NODE_CLASS_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(ASYNC_TEST_NODE_CLASS_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(API_TEST_NODE_CLASS_MAPPINGS)
|
||||
NODE_CLASS_MAPPINGS.update(CACHE_PROVIDER_TEST_NODE_CLASS_MAPPINGS)
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(TEST_NODE_DISPLAY_NAME_MAPPINGS)
|
||||
@@ -28,4 +26,3 @@ NODE_DISPLAY_NAME_MAPPINGS.update(CONDITION_NODE_DISPLAY_NAME_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(TEST_STUB_NODE_DISPLAY_NAME_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(ASYNC_TEST_NODE_DISPLAY_NAME_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(API_TEST_NODE_DISPLAY_NAME_MAPPINGS)
|
||||
NODE_DISPLAY_NAME_MAPPINGS.update(CACHE_PROVIDER_TEST_NODE_DISPLAY_NAME_MAPPINGS)
|
||||
|
||||
@@ -135,148 +135,6 @@ class TestSyncError(ComfyNodeABC):
|
||||
raise RuntimeError("Intentional sync execution error for testing")
|
||||
|
||||
|
||||
class TestOOMError(ComfyNodeABC):
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"value": (IO.ANY, {})}}
|
||||
|
||||
RETURN_TYPES = (IO.ANY,)
|
||||
FUNCTION = "oom_error"
|
||||
CATEGORY = "experimental/async"
|
||||
|
||||
def oom_error(self, value):
|
||||
raise torch.OutOfMemoryError("Intentional out of memory error for testing")
|
||||
|
||||
|
||||
class TestMixedExpansionFailure(ComfyNodeABC):
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"value": ("INT", {})}}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "expand"
|
||||
CATEGORY = "experimental/async"
|
||||
|
||||
def expand(self, value):
|
||||
image = torch.zeros([1, 32, 32, 3])
|
||||
if value == 0:
|
||||
return (image,)
|
||||
|
||||
graph = GraphBuilder()
|
||||
error = graph.node("TestSyncError", value=image)
|
||||
return {
|
||||
"result": (error.out(0),),
|
||||
"expand": graph.finalize(),
|
||||
}
|
||||
|
||||
|
||||
class TestMalformedExpansion(ComfyNodeABC):
|
||||
"""Expands to a graph referencing a missing node class, so the failure
|
||||
happens in the executor after the node function has returned."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"value": (IO.ANY, {})}}
|
||||
|
||||
RETURN_TYPES = (IO.ANY,)
|
||||
FUNCTION = "expand"
|
||||
CATEGORY = "experimental/async"
|
||||
|
||||
def expand(self, value):
|
||||
graph = GraphBuilder()
|
||||
missing = graph.node("TestNodeClassThatDoesNotExist", value=value)
|
||||
return {
|
||||
"result": (missing.out(0),),
|
||||
"expand": graph.finalize(),
|
||||
}
|
||||
|
||||
|
||||
class TestMalformedResult(ComfyNodeABC):
|
||||
"""Returns a non-tuple result so the failure happens while the executor
|
||||
merges results, after the node function has returned."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"value": (IO.ANY, {})}}
|
||||
|
||||
RETURN_TYPES = (IO.ANY,)
|
||||
FUNCTION = "run"
|
||||
CATEGORY = "experimental/async"
|
||||
|
||||
def run(self, value):
|
||||
return 5
|
||||
|
||||
|
||||
class TestCyclicExpansion(ComfyNodeABC):
|
||||
"""Expands to an output node that consumes this node's own pending output,
|
||||
forming a cycle through the expansion completion link."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {"value": (IO.ANY, {})},
|
||||
"hidden": {"unique_id": "UNIQUE_ID"},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (IO.ANY,)
|
||||
FUNCTION = "expand"
|
||||
CATEGORY = "experimental/async"
|
||||
|
||||
def expand(self, value, unique_id):
|
||||
graph = GraphBuilder()
|
||||
graph.node("TestAsyncOutput", value=[unique_id, 0], seconds=0.0)
|
||||
return {
|
||||
"result": (value,),
|
||||
"expand": graph.finalize(),
|
||||
}
|
||||
|
||||
|
||||
class TestExpansionWithFailingOutput(ComfyNodeABC):
|
||||
"""Expands to a subgraph whose result succeeds while a side branch ending
|
||||
in an output node fails."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"image": (IO.IMAGE, {})}}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "expand"
|
||||
CATEGORY = "experimental/async"
|
||||
|
||||
def expand(self, image):
|
||||
graph = GraphBuilder()
|
||||
error = graph.node("TestSyncError", value=image)
|
||||
graph.node("PreviewImage", images=error.out(0))
|
||||
passthrough = graph.node("StubImage", content="WHITE", height=32, width=32, batch_size=1)
|
||||
return {
|
||||
"result": (passthrough.out(0),),
|
||||
"expand": graph.finalize(),
|
||||
}
|
||||
|
||||
|
||||
class TestAsyncOutput(ComfyNodeABC):
|
||||
"""Async output node with no return sockets, used to test partial failure
|
||||
handling across pending async invocations."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"value": (IO.ANY, {}),
|
||||
"seconds": (IO.FLOAT, {"default": 0.1}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ()
|
||||
OUTPUT_NODE = True
|
||||
FUNCTION = "run"
|
||||
CATEGORY = "experimental/async"
|
||||
|
||||
async def run(self, value, seconds=0.1):
|
||||
await asyncio.sleep(seconds)
|
||||
return {"ui": {"values": [1]}}
|
||||
|
||||
|
||||
class TestAsyncLazyCheck(ComfyNodeABC):
|
||||
"""Test node with async check_lazy_status."""
|
||||
|
||||
@@ -464,13 +322,6 @@ ASYNC_TEST_NODE_CLASS_MAPPINGS = {
|
||||
"TestAsyncValidationError": TestAsyncValidationError,
|
||||
"TestAsyncTimeout": TestAsyncTimeout,
|
||||
"TestSyncError": TestSyncError,
|
||||
"TestOOMError": TestOOMError,
|
||||
"TestMixedExpansionFailure": TestMixedExpansionFailure,
|
||||
"TestMalformedExpansion": TestMalformedExpansion,
|
||||
"TestMalformedResult": TestMalformedResult,
|
||||
"TestCyclicExpansion": TestCyclicExpansion,
|
||||
"TestExpansionWithFailingOutput": TestExpansionWithFailingOutput,
|
||||
"TestAsyncOutput": TestAsyncOutput,
|
||||
"TestAsyncLazyCheck": TestAsyncLazyCheck,
|
||||
"TestDynamicAsyncGeneration": TestDynamicAsyncGeneration,
|
||||
"TestAsyncResourceUser": TestAsyncResourceUser,
|
||||
@@ -484,13 +335,6 @@ ASYNC_TEST_NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"TestAsyncValidationError": "Test Async Validation Error",
|
||||
"TestAsyncTimeout": "Test Async Timeout",
|
||||
"TestSyncError": "Test Sync Error",
|
||||
"TestOOMError": "Test OOM Error",
|
||||
"TestMixedExpansionFailure": "Test Mixed Expansion Failure",
|
||||
"TestMalformedExpansion": "Test Malformed Expansion",
|
||||
"TestMalformedResult": "Test Malformed Result",
|
||||
"TestCyclicExpansion": "Test Cyclic Expansion",
|
||||
"TestExpansionWithFailingOutput": "Test Expansion With Failing Output",
|
||||
"TestAsyncOutput": "Test Async Output",
|
||||
"TestAsyncLazyCheck": "Test Async Lazy Check",
|
||||
"TestDynamicAsyncGeneration": "Test Dynamic Async Generation",
|
||||
"TestAsyncResourceUser": "Test Async Resource User",
|
||||
|
||||
@@ -1,51 +0,0 @@
|
||||
from comfy.comfy_types.node_typing import ComfyNodeABC
|
||||
from comfy_api.latest._caching import CacheProvider
|
||||
from comfy_execution.cache_provider import register_cache_provider
|
||||
|
||||
|
||||
class _RecordingCacheProvider(CacheProvider):
|
||||
"""Records the class types of every externally stored cache entry so tests
|
||||
can assert that failed or failure-blocked outputs never leave the process."""
|
||||
|
||||
def __init__(self):
|
||||
self.stored_class_types = []
|
||||
|
||||
async def on_lookup(self, context):
|
||||
return None
|
||||
|
||||
async def on_store(self, context, value):
|
||||
self.stored_class_types.append(context.class_type)
|
||||
|
||||
|
||||
RECORDING_CACHE_PROVIDER = _RecordingCacheProvider()
|
||||
register_cache_provider(RECORDING_CACHE_PROVIDER)
|
||||
|
||||
|
||||
class TestCacheProviderRecord(ComfyNodeABC):
|
||||
"""Reports which node class types have been stored through the external
|
||||
cache provider interface since the server started."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {}}
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(cls):
|
||||
return float("NaN")
|
||||
|
||||
RETURN_TYPES = ()
|
||||
OUTPUT_NODE = True
|
||||
FUNCTION = "report"
|
||||
CATEGORY = "Testing/Nodes"
|
||||
|
||||
def report(self):
|
||||
return {"ui": {"stored_class_types": list(RECORDING_CACHE_PROVIDER.stored_class_types)}}
|
||||
|
||||
|
||||
CACHE_PROVIDER_TEST_NODE_CLASS_MAPPINGS = {
|
||||
"TestCacheProviderRecord": TestCacheProviderRecord,
|
||||
}
|
||||
|
||||
CACHE_PROVIDER_TEST_NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"TestCacheProviderRecord": "Test Cache Provider Record",
|
||||
}
|
||||
Reference in New Issue
Block a user