mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-14 07:46:39 +08:00
Compare commits
12 Commits
feat/parti
...
master
| Author | SHA1 | Date | |
|---|---|---|---|
| 80acfcf0fe | |||
| c35a622acd | |||
| da2608926e | |||
| 5658a68a87 | |||
| 5bb831a3f5 | |||
| 5697b97017 | |||
| ec0e8b3447 | |||
| 8deaa4d911 | |||
| b58f829b57 | |||
| 917faef771 | |||
| 8b099de36a | |||
| 69ea58697b |
@ -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, **kwargs):
|
||||
def two_pass_attention(q, k, v, heads, enable_gqa=False, **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)
|
||||
out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options, enable_gqa=enable_gqa)
|
||||
elif ar_len >= T:
|
||||
out = comfy.ops.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True)
|
||||
out = comfy.ops.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True, enable_gqa=enable_gqa)
|
||||
elif ar_len <= 0:
|
||||
out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options)
|
||||
out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options, enable_gqa=enable_gqa)
|
||||
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,
|
||||
attn_mask=None, dropout_p=0.0, is_causal=True, enable_gqa=enable_gqa,
|
||||
)
|
||||
out_gen = optimized_attention(
|
||||
q[:, :, ar_len:], k, v, heads,
|
||||
mask=None, skip_reshape=True, skip_output_reshape=True,
|
||||
transformer_options=transformer_options,
|
||||
transformer_options=transformer_options, enable_gqa=enable_gqa,
|
||||
)
|
||||
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("flash_attention::flash_attn", mutates_args=())
|
||||
@torch.library.custom_op("comfy::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
|
||||
|
||||
@ -197,6 +197,9 @@ 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))
|
||||
@ -226,6 +229,7 @@ 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:
|
||||
|
||||
@ -13,15 +13,15 @@ from .model import PixDiT_T2I
|
||||
from .modules import precompute_freqs_cis_2d
|
||||
|
||||
|
||||
class SigmaAwareGatePerTokenPerDim(nn.Module):
|
||||
class SigmaAwareGate(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, dtype=None, device=None, operations=None):
|
||||
def __init__(self, dim: int, per_token: bool = False, dtype=None, device=None, operations=None):
|
||||
super().__init__()
|
||||
self.content_proj = operations.Linear(dim * 2, dim, dtype=dtype, device=device)
|
||||
self.content_proj = operations.Linear(dim * 2, 1 if per_token else 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 SigmaAwareGatePerTokenPerDim(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, dtype=None, device=None, operations=None):
|
||||
def __init__(self, channels: int, num_groups: int = 4, conv_padding_mode: str = "zeros", 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, dtype=dtype, device=device),
|
||||
operations.Conv2d(channels, channels, kernel_size=3, padding=1, padding_mode=conv_padding_mode, dtype=dtype, device=device),
|
||||
operations.GroupNorm(num_groups, channels, dtype=dtype, device=device),
|
||||
nn.SiLU(),
|
||||
operations.Conv2d(channels, channels, kernel_size=3, padding=1, dtype=dtype, device=device),
|
||||
operations.Conv2d(channels, channels, kernel_size=3, padding=1, padding_mode=conv_padding_mode, dtype=dtype, device=device),
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
@ -62,9 +62,13 @@ 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__()
|
||||
@ -74,34 +78,38 @@ 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
|
||||
|
||||
z_to_patch_ratio = (sr_scale * latent_spatial_down_factor) / patch_size
|
||||
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
|
||||
self.z_to_patch_ratio = z_to_patch_ratio
|
||||
if z_to_patch_ratio >= 1:
|
||||
self.latent_fold_factor = 0
|
||||
latent_proj_in_ch = latent_channels
|
||||
latent_proj_in_ch = effective_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 = latent_channels * fold_factor * fold_factor
|
||||
latent_proj_in_ch = effective_latent_channels * fold_factor * fold_factor
|
||||
|
||||
layers = [
|
||||
operations.Conv2d(latent_proj_in_ch, hidden_dim, kernel_size=3, padding=1, dtype=dtype, device=device),
|
||||
operations.Conv2d(latent_proj_in_ch, hidden_dim, kernel_size=3, padding=1, padding_mode=conv_padding_mode, dtype=dtype, device=device),
|
||||
nn.SiLU(),
|
||||
operations.Conv2d(hidden_dim, hidden_dim, kernel_size=3, padding=1, dtype=dtype, device=device),
|
||||
operations.Conv2d(hidden_dim, hidden_dim, kernel_size=3, padding=1, padding_mode=conv_padding_mode, dtype=dtype, device=device),
|
||||
]
|
||||
for _ in range(num_res_blocks):
|
||||
layers.append(ResBlock(hidden_dim, dtype=dtype, device=device, operations=operations))
|
||||
layers.append(ResBlock(hidden_dim, conv_padding_mode=conv_padding_mode, 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(
|
||||
[SigmaAwareGatePerTokenPerDim(out_dim, dtype=dtype, device=device, operations=operations)
|
||||
[SigmaAwareGate(out_dim, per_token=gate_per_token, dtype=dtype, device=device, operations=operations)
|
||||
for _ in range(num_outputs)]
|
||||
)
|
||||
|
||||
@ -115,6 +123,11 @@ 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:
|
||||
@ -134,7 +147,10 @@ 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)
|
||||
return [head(tokens) for head in self.output_heads]
|
||||
outputs = [head(tokens) for head in self.output_heads]
|
||||
if self.pit_head is not None:
|
||||
outputs.append(self.pit_head(tokens))
|
||||
return outputs
|
||||
|
||||
|
||||
class PidNet(PixDiT_T2I):
|
||||
@ -148,6 +164,10 @@ 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,
|
||||
@ -165,6 +185,8 @@ 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,
|
||||
@ -173,13 +195,20 @@ 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(
|
||||
@ -197,6 +226,11 @@ 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")
|
||||
@ -216,12 +250,14 @@ 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,
|
||||
)
|
||||
|
||||
@ -470,15 +470,46 @@ 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:
|
||||
in_ch = int(state_dict[_lq_w_key].shape[1])
|
||||
latent_proj_in_channels = int(state_dict[_lq_w_key].shape[1])
|
||||
hidden_dim = int(state_dict[_lq_w_key].shape[0])
|
||||
_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_latent_channels": in_ch,
|
||||
"latent_spatial_down_factor": 16 if in_ch >= 64 else 8}
|
||||
"lq_hidden_dim": hidden_dim}
|
||||
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
|
||||
|
||||
@ -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}},
|
||||
|
||||
@ -1133,7 +1133,9 @@ 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-preview"
|
||||
model = "gemini-3.1-flash-image"
|
||||
elif model == "gemini-3-pro-image-preview":
|
||||
model = "gemini-3-pro-image"
|
||||
|
||||
parts: list[GeminiPart] = [GeminiPart(text=prompt)]
|
||||
if images is not None:
|
||||
@ -1507,7 +1509,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-preview"
|
||||
model_id = "gemini-3.1-flash-image"
|
||||
elif model_choice == "Nano Banana 2 Lite":
|
||||
model_id = "gemini-3.1-flash-lite-image"
|
||||
else:
|
||||
|
||||
@ -15,6 +15,7 @@ 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
|
||||
|
||||
@ -60,6 +61,7 @@ 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,
|
||||
}
|
||||
|
||||
|
||||
|
||||
@ -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,27 +201,19 @@ class ExecutionList(TopologicalSort):
|
||||
self.staged_node_id = None
|
||||
self.execution_cache = {}
|
||||
self.execution_cache_listeners = {}
|
||||
self.transient_cache = {}
|
||||
|
||||
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)
|
||||
@ -232,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:
|
||||
|
||||
@ -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
|
||||
|
||||
@ -204,10 +204,8 @@ 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
|
||||
if status_info:
|
||||
messages = status_info.get('messages', [])
|
||||
@ -219,23 +217,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)
|
||||
|
||||
execution_summary = (status_info.get('execution_summary') if status_info else None) or {}
|
||||
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
|
||||
@ -252,9 +237,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,
|
||||
@ -263,7 +245,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,
|
||||
|
||||
@ -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():
|
||||
|
||||
@ -844,15 +844,18 @@ class ImageMergeTileList(IO.ComfyNode):
|
||||
# Format specifications
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Maps (file_format, bit_depth, has_alpha) -> (numpy dtype scale, av pixel format,
|
||||
# stream pix_fmt). Keeps the encode path declarative instead of branchy.
|
||||
# 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.
|
||||
_FORMAT_SPECS = {
|
||||
("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"},
|
||||
("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"},
|
||||
}
|
||||
|
||||
|
||||
@ -891,10 +894,11 @@ 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 inside the log branch so negative / out-of-range values don't blow up;
|
||||
# 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;
|
||||
# values above 1.0 are allowed and extrapolate naturally.
|
||||
low = (t ** 2) / 3.0
|
||||
high = (torch.exp((t.clamp(min=_HLG_C) - _HLG_C) / _HLG_A) + _HLG_B) / 12.0
|
||||
high = (torch.exp((t.clamp(min=0.5) - _HLG_C) / _HLG_A) + _HLG_B) / 12.0
|
||||
return torch.where(t <= 0.5, low, high)
|
||||
|
||||
|
||||
@ -1087,7 +1091,8 @@ def _encode_image(
|
||||
bit_depth: str,
|
||||
colorspace: str,
|
||||
) -> bytes:
|
||||
"""Encode a single HxWxC tensor to PNG or EXR bytes in memory.
|
||||
"""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.
|
||||
|
||||
For EXR the input is interpreted according to `colorspace` and converted
|
||||
to scene-linear (EXR's convention) before writing:
|
||||
@ -1101,10 +1106,16 @@ 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[(file_format, bit_depth, has_alpha)]
|
||||
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)."
|
||||
)
|
||||
|
||||
if spec["dtype"] == np.float32:
|
||||
# EXR path: preserve full range, no clamp.
|
||||
|
||||
@ -81,7 +81,7 @@ class SaveVideo(io.ComfyNode):
|
||||
display_name="Save Video",
|
||||
category="video",
|
||||
essentials_category="Basics",
|
||||
description="Saves the input images to your ComfyUI output directory.",
|
||||
description="Saves the input videos 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."),
|
||||
|
||||
167
execution.py
167
execution.py
@ -32,13 +32,9 @@ from comfy_execution.caching import (
|
||||
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 +51,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,38 +104,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 _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
|
||||
@ -198,7 +152,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)
|
||||
@ -495,8 +449,6 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
|
||||
return (ExecutionResult.SUCCESS, None, None)
|
||||
|
||||
input_data_all = None
|
||||
failure_blocked_invocations = 0
|
||||
successful_invocations = 0
|
||||
try:
|
||||
if unique_id in pending_async_nodes:
|
||||
results = []
|
||||
@ -566,9 +518,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,
|
||||
@ -587,8 +536,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)
|
||||
|
||||
@ -664,14 +611,8 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
|
||||
return (ExecutionResult.PENDING, None, None)
|
||||
|
||||
cache_entry = CacheEntry(ui=ui_outputs.get(unique_id), outputs=output_data)
|
||||
if extra_data.get(NODE_FAILURE_POLICY_EXTRA_DATA_KEY) == NODE_FAILURE_POLICY_CONTINUE_INDEPENDENT:
|
||||
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
|
||||
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")
|
||||
@ -713,12 +654,6 @@ async def execute(server, dynprompt, caches, current_item, extra_data, executed,
|
||||
|
||||
return (ExecutionResult.FAILURE, error_details, ex)
|
||||
|
||||
fully_failure_blocked = failure_blocked_invocations > 0 and successful_invocations == 0
|
||||
fully_failure_blocked = fully_failure_blocked 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)
|
||||
|
||||
@ -735,7 +670,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 = {}
|
||||
|
||||
def add_message(self, event, data: dict, broadcast: bool):
|
||||
data = {
|
||||
@ -774,21 +708,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
|
||||
@ -809,8 +728,6 @@ class PromptExecutor:
|
||||
set_preview_method(extra_data.get("preview_method"))
|
||||
|
||||
nodes.interrupt_processing(False)
|
||||
self.success = True
|
||||
self.execution_summary = {}
|
||||
|
||||
if "client_id" in extra_data:
|
||||
self.server.client_id = extra_data["client_id"]
|
||||
@ -855,59 +772,24 @@ 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 _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)
|
||||
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:
|
||||
@ -935,41 +817,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.execution_summary = {
|
||||
"completion_status": "success",
|
||||
"has_errors": False,
|
||||
"execution_error_count": 0,
|
||||
}
|
||||
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 = {}
|
||||
@ -1427,7 +1275,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):
|
||||
|
||||
3
main.py
3
main.py
@ -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)
|
||||
|
||||
|
||||
52
openapi.yaml
52
openapi.yaml
@ -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,6 +144,14 @@ 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}$
|
||||
@ -775,14 +783,6 @@ components:
|
||||
ModelFolder:
|
||||
description: Represents a folder containing models
|
||||
properties:
|
||||
extensions:
|
||||
description: The folder's registered file-extension allowlist. An empty array means the folder accepts any extension (match-all).
|
||||
example:
|
||||
- .ckpt
|
||||
- .safetensors
|
||||
items:
|
||||
type: string
|
||||
type: array
|
||||
folders:
|
||||
description: List of paths where models of this type are stored
|
||||
example:
|
||||
@ -922,13 +922,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 +1644,7 @@ paths:
|
||||
format: uuid
|
||||
type: string
|
||||
tags:
|
||||
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.
|
||||
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.
|
||||
type: string
|
||||
user_metadata:
|
||||
description: Custom JSON metadata as a string
|
||||
@ -1836,7 +1829,7 @@ paths:
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/Asset'
|
||||
$ref: '#/components/schemas/AssetUpdated'
|
||||
description: Asset updated successfully
|
||||
"400":
|
||||
content:
|
||||
@ -2477,9 +2470,6 @@ 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:
|
||||
@ -3307,6 +3297,12 @@ 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:
|
||||
|
||||
@ -1,6 +1,6 @@
|
||||
comfyui-frontend-package==1.45.20
|
||||
comfyui-workflow-templates==0.11.6
|
||||
comfyui-embedded-docs==0.5.7
|
||||
comfyui-workflow-templates==0.11.9
|
||||
comfyui-embedded-docs==0.5.8
|
||||
torch
|
||||
torchsde
|
||||
torchvision
|
||||
@ -22,7 +22,7 @@ alembic
|
||||
SQLAlchemy>=2.0.0
|
||||
filelock
|
||||
av>=16.0.0
|
||||
comfy-kitchen==0.2.18
|
||||
comfy-kitchen==0.2.19
|
||||
comfy-aimdo==0.4.10
|
||||
requests
|
||||
simpleeval>=1.0.0
|
||||
|
||||
17
server.py
17
server.py
@ -1097,19 +1097,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 +1104,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"]
|
||||
|
||||
|
||||
@ -97,6 +97,21 @@ def _make_seedvr2_3b_shared_mm_sd():
|
||||
}
|
||||
|
||||
|
||||
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()}
|
||||
|
||||
@ -206,6 +221,43 @@ class TestModelDetection:
|
||||
|
||||
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
|
||||
|
||||
@ -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):
|
||||
|
||||
@ -13,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)
|
||||
@ -55,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)
|
||||
@ -96,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())
|
||||
@ -165,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()
|
||||
@ -184,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', [])
|
||||
@ -344,122 +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_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 = {
|
||||
|
||||
@ -135,41 +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 TestAsyncLazyCheck(ComfyNodeABC):
|
||||
"""Test node with async check_lazy_status."""
|
||||
|
||||
@ -357,8 +322,6 @@ ASYNC_TEST_NODE_CLASS_MAPPINGS = {
|
||||
"TestAsyncValidationError": TestAsyncValidationError,
|
||||
"TestAsyncTimeout": TestAsyncTimeout,
|
||||
"TestSyncError": TestSyncError,
|
||||
"TestOOMError": TestOOMError,
|
||||
"TestMixedExpansionFailure": TestMixedExpansionFailure,
|
||||
"TestAsyncLazyCheck": TestAsyncLazyCheck,
|
||||
"TestDynamicAsyncGeneration": TestDynamicAsyncGeneration,
|
||||
"TestAsyncResourceUser": TestAsyncResourceUser,
|
||||
@ -372,8 +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",
|
||||
"TestAsyncLazyCheck": "Test Async Lazy Check",
|
||||
"TestDynamicAsyncGeneration": "Test Dynamic Async Generation",
|
||||
"TestAsyncResourceUser": "Test Async Resource User",
|
||||
|
||||
Reference in New Issue
Block a user