Merge branch 'master' into fix/core/video-transcode-oom
Some checks failed
Python Linting / Run Ruff (push) Has been cancelled
Python Linting / Run Pylint (push) Has been cancelled
Build package / Build Test (3.10) (push) Has been cancelled
Build package / Build Test (3.11) (push) Has been cancelled
Build package / Build Test (3.12) (push) Has been cancelled
Build package / Build Test (3.13) (push) Has been cancelled
Build package / Build Test (3.14) (push) Has been cancelled

This commit is contained in:
Alexis Rolland 2026-07-15 15:19:26 +08:00 committed by GitHub
commit 9b53498f9d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
64 changed files with 7061 additions and 142 deletions

View File

@ -32,9 +32,11 @@ jobs:
PR_NUMBER: ${{ github.event.pull_request.number || github.event.issue.number }} PR_NUMBER: ${{ github.event.pull_request.number || github.event.issue.number }}
PR_AUTHOR: ${{ github.event.pull_request.user.login || github.event.issue.user.login }} 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] 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: | run: |
others=$(gh api "repos/${{ github.repository }}/pulls/${PR_NUMBER}/commits" --paginate \ others=$(gh api "repos/${{ github.repository }}/pulls/${PR_NUMBER}/commits" --paginate \
--jq '.[] | (.author.login // empty), (.committer.login // empty)' \ --jq '.[] | (.author.login // .commit.author.name // empty), (.committer.login // .commit.committer.name // empty)' \
| sort -u | grep -vix "${PR_AUTHOR}" | paste -sd, -) | sort -u | grep -vix "${PR_AUTHOR}" | paste -sd, -)
if [ -n "$others" ]; then if [ -n "$others" ]; then
echo "allowlist=${BASE_ALLOWLIST},${others}" >> "$GITHUB_OUTPUT" echo "allowlist=${BASE_ALLOWLIST},${others}" >> "$GITHUB_OUTPUT"
@ -43,7 +45,7 @@ jobs:
fi fi
- name: CLA Assistant - name: CLA Assistant
# Run on PR events, on "recheck" comment, or when someone posts the exact signing phrase. # Run on PR events, on "recheck" comment, or when someone posts the signing phrase.
# IMPORTANT: this phrase must match `custom-pr-sign-comment` below. # IMPORTANT: this phrase must match `custom-pr-sign-comment` below.
if: > if: >
github.event_name == 'pull_request_target' || github.event_name == 'pull_request_target' ||

View File

@ -35,7 +35,11 @@ class ModelFileManager:
for folder in model_types: for folder in model_types:
if folder in folder_black_list: if folder in folder_black_list:
continue continue
output_folders.append({"name": folder, "folders": folder_paths.get_folder_paths(folder)}) output_folders.append({
"name": folder,
"folders": folder_paths.get_folder_paths(folder),
"extensions": sorted(folder_paths.folder_names_and_paths[folder][1]),
})
return web.json_response(output_folders) return web.json_response(output_folders)
# NOTE: This is an experiment to replace `/models/{folder}` # NOTE: This is an experiment to replace `/models/{folder}`

View File

@ -92,6 +92,7 @@ parser.add_argument("--directml", type=int, nargs="?", metavar="DIRECTML_DEVICE"
parser.add_argument("--oneapi-device-selector", type=str, default=None, metavar="SELECTOR_STRING", help="Sets the oneAPI device(s) this instance will use.") parser.add_argument("--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("--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("--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): class LatentPreviewMethod(enum.Enum):
NoPreviews = "none" NoPreviews = "none"

View File

@ -779,6 +779,10 @@ class ACEAudio(LatentFormat):
latent_channels = 8 latent_channels = 8
latent_dimensions = 2 latent_dimensions = 2
class SeedVR2(LatentFormat):
latent_channels = 16
latent_dimensions = 3
class ACEAudio15(LatentFormat): class ACEAudio15(LatentFormat):
latent_channels = 64 latent_channels = 64
latent_dimensions = 1 latent_dimensions = 1

View File

@ -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. 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 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 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: 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: 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: else:
out_ar = comfy.ops.scaled_dot_product_attention( out_ar = comfy.ops.scaled_dot_product_attention(
q[:, :, :ar_len], k[:, :, :ar_len], v[:, :, :ar_len], 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( out_gen = optimized_attention(
q[:, :, ar_len:], k, v, heads, q[:, :, ar_len:], k, v, heads,
mask=None, skip_reshape=True, skip_output_reshape=True, 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) out = torch.cat([out_ar, out_gen], dim=2)

View File

@ -709,7 +709,7 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
return out return out
try: 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, 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: 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 softmax_scale_arg = None if softmax_scale == -1.0 else softmax_scale

View File

@ -22,7 +22,7 @@ def torch_cat_if_needed(xl, dim):
else: else:
return None return None
def get_timestep_embedding(timesteps, embedding_dim): def get_timestep_embedding(timesteps, embedding_dim, flip_sin_to_cos=False, downscale_freq_shift=1):
""" """
This matches the implementation in Denoising Diffusion Probabilistic Models: This matches the implementation in Denoising Diffusion Probabilistic Models:
From Fairseq. From Fairseq.
@ -33,11 +33,13 @@ def get_timestep_embedding(timesteps, embedding_dim):
assert len(timesteps.shape) == 1 assert len(timesteps.shape) == 1
half_dim = embedding_dim // 2 half_dim = embedding_dim // 2
emb = math.log(10000) / (half_dim - 1) emb = math.log(10000) / (half_dim - downscale_freq_shift)
emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb) emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb)
emb = emb.to(device=timesteps.device) emb = emb.to(device=timesteps.device)
emb = timesteps.float()[:, None] * emb[None, :] emb = timesteps.float()[:, None] * emb[None, :]
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1) 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 if embedding_dim % 2 == 1: # zero pad
emb = torch.nn.functional.pad(emb, (0,1,0,0)) emb = torch.nn.functional.pad(emb, (0,1,0,0))
return emb return emb

View File

@ -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).""" """Hook for subclasses to inject per-block state into the patch stream (e.g. PiD's LQ gate)."""
return s return s
def _pre_pixel_blocks(self, s, **kwargs):
return s
def _forward(self, x, timesteps, context=None, attention_mask=None, transformer_options={}, **kwargs): def _forward(self, x, timesteps, context=None, attention_mask=None, transformer_options={}, **kwargs):
H_orig, W_orig = x.shape[2], x.shape[3] 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)) 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, y_emb = blk(s, y_emb, condition, pos_img, pos_txt, None, transformer_options=transformer_options)
s = F.silu(t_emb + s) s = F.silu(t_emb + s)
s = self._pre_pixel_blocks(s, **kwargs)
s_cond = s.view(B * L, self.hidden_size) s_cond = s.view(B * L, self.hidden_size)
x_pixels = self.pixel_embedder(x, patch_size=self.patch_size) x_pixels = self.pixel_embedder(x, patch_size=self.patch_size)
for blk in self.pixel_blocks: for blk in self.pixel_blocks:

View File

@ -13,15 +13,15 @@ from .model import PixDiT_T2I
from .modules import precompute_freqs_cis_2d 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. """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. 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__() 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)) 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: 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): class ResBlock(nn.Module):
"""Pre-activation ResNet block: GN -> SiLU -> Conv -> GN -> SiLU -> Conv + skip.""" """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__() super().__init__()
self.block = nn.Sequential( self.block = nn.Sequential(
operations.GroupNorm(num_groups, channels, dtype=dtype, device=device), operations.GroupNorm(num_groups, channels, dtype=dtype, device=device),
nn.SiLU(), 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), operations.GroupNorm(num_groups, channels, dtype=dtype, device=device),
nn.SiLU(), 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: def forward(self, x: torch.Tensor) -> torch.Tensor:
@ -62,9 +62,13 @@ class LQProjection2D(nn.Module):
patch_size: int = 16, patch_size: int = 16,
sr_scale: int = 4, sr_scale: int = 4,
latent_spatial_down_factor: int = 8, latent_spatial_down_factor: int = 8,
latent_unpatchify_factor: int = 1,
num_res_blocks: int = 4, num_res_blocks: int = 4,
num_outputs: int = 7, num_outputs: int = 7,
interval: int = 2, interval: int = 2,
conv_padding_mode: str = "zeros",
gate_per_token: bool = False,
pit_output: bool = False,
dtype=None, device=None, operations=None, dtype=None, device=None, operations=None,
): ):
super().__init__() super().__init__()
@ -74,34 +78,38 @@ class LQProjection2D(nn.Module):
self.patch_size = patch_size self.patch_size = patch_size
self.sr_scale = sr_scale self.sr_scale = sr_scale
self.latent_spatial_down_factor = latent_spatial_down_factor self.latent_spatial_down_factor = latent_spatial_down_factor
self.latent_unpatchify_factor = latent_unpatchify_factor
self.num_outputs = num_outputs self.num_outputs = num_outputs
self.interval = interval 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 self.z_to_patch_ratio = z_to_patch_ratio
if z_to_patch_ratio >= 1: if z_to_patch_ratio >= 1:
self.latent_fold_factor = 0 self.latent_fold_factor = 0
latent_proj_in_ch = latent_channels latent_proj_in_ch = effective_latent_channels
else: else:
fold_factor = int(1 / z_to_patch_ratio) fold_factor = int(1 / z_to_patch_ratio)
assert fold_factor * z_to_patch_ratio == 1.0 assert fold_factor * z_to_patch_ratio == 1.0
self.latent_fold_factor = fold_factor 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 = [ 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(), 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): 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.latent_proj = nn.Sequential(*layers)
self.output_heads = nn.ModuleList( self.output_heads = nn.ModuleList(
[operations.Linear(hidden_dim, out_dim, dtype=dtype, device=device) for _ in range(num_outputs)] [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( 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)] for _ in range(num_outputs)]
) )
@ -115,6 +123,11 @@ class LQProjection2D(nn.Module):
return self.gate_modules[out_idx](x, lq_feature, sigma) 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: 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] B, z_dim = lq_latent.shape[:2]
if self.z_to_patch_ratio >= 1: if self.z_to_patch_ratio >= 1:
if lq_latent.shape[2] != pH or lq_latent.shape[3] != pW: 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) feat = self._align_latent_to_patch_grid(lq_latent, target_pH, target_pW)
B, C, H, W = feat.shape B, C, H, W = feat.shape
tokens = feat.permute(0, 2, 3, 1).contiguous().view(B, H * W, C) 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): class PidNet(PixDiT_T2I):
@ -148,6 +164,10 @@ class PidNet(PixDiT_T2I):
lq_interval: int = 2, lq_interval: int = 2,
sr_scale: int = 4, sr_scale: int = 4,
latent_spatial_down_factor: int = 8, 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_h: int = 1024, # NTK ref resolution in PIXEL units: 1024px / patch=16 -> grid_ref=64.
rope_ref_w: int = 1024, rope_ref_w: int = 1024,
image_model=None, image_model=None,
@ -165,6 +185,8 @@ class PidNet(PixDiT_T2I):
for blk in self.pixel_blocks: for blk in self.pixel_blocks:
blk._rope_fn = _pit_rope_fn blk._rope_fn = _pit_rope_fn
self.pit_lq_inject = pit_lq_inject
num_lq_outputs = (self.patch_depth + lq_interval - 1) // lq_interval num_lq_outputs = (self.patch_depth + lq_interval - 1) // lq_interval
self.lq_proj = LQProjection2D( self.lq_proj = LQProjection2D(
latent_channels=lq_latent_channels, latent_channels=lq_latent_channels,
@ -173,13 +195,20 @@ class PidNet(PixDiT_T2I):
patch_size=self.patch_size, patch_size=self.patch_size,
sr_scale=sr_scale, sr_scale=sr_scale,
latent_spatial_down_factor=latent_spatial_down_factor, latent_spatial_down_factor=latent_spatial_down_factor,
latent_unpatchify_factor=lq_latent_unpatchify_factor,
num_res_blocks=lq_num_res_blocks, num_res_blocks=lq_num_res_blocks,
num_outputs=num_lq_outputs, num_outputs=num_lq_outputs,
interval=lq_interval, interval=lq_interval,
conv_padding_mode=lq_conv_padding_mode,
gate_per_token=lq_gate_per_token,
pit_output=pit_lq_inject,
dtype=dtype, dtype=dtype,
device=device, device=device,
operations=operations, 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): def _fetch_patch_pos(self, height, width, device, dtype, **rope_opts):
return precompute_freqs_cis_2d( return precompute_freqs_cis_2d(
@ -197,6 +226,11 @@ class PidNet(PixDiT_T2I):
return s return s
return self.lq_proj.gate(s, pid_lq_features[out_idx], pid_degrade_sigma, out_idx) 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): def _forward(self, x, timesteps, context=None, attention_mask=None, transformer_options={}, lq_latent=None, degrade_sigma=None, **kwargs):
if lq_latent is None: if lq_latent is None:
raise ValueError("PidNet requires lq_latent — attach via PiDConditioning") raise ValueError("PidNet requires lq_latent — attach via PiDConditioning")
@ -216,12 +250,14 @@ class PidNet(PixDiT_T2I):
degrade_sigma = degrade_sigma.expand(B).contiguous() degrade_sigma = degrade_sigma.expand(B).contiguous()
lq_features = self.lq_proj(lq_latent=lq_latent.to(x), target_pH=Hs, target_pW=Ws) 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( return super()._forward(
x, timesteps, x, timesteps,
context=context, attention_mask=attention_mask, context=context, attention_mask=attention_mask,
transformer_options=transformer_options, transformer_options=transformer_options,
pid_lq_features=lq_features, pid_lq_features=lq_features,
pid_pit_lq_feature=pit_lq_feature,
pid_degrade_sigma=degrade_sigma, pid_degrade_sigma=degrade_sigma,
**kwargs, **kwargs,
) )

View File

@ -0,0 +1,51 @@
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

View File

@ -0,0 +1,301 @@
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

View File

@ -0,0 +1,48 @@
"""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).

1361
comfy/ldm/seedvr/model.py Normal file

File diff suppressed because it is too large Load Diff

1610
comfy/ldm/seedvr/vae.py Normal file

File diff suppressed because it is too large Load Diff

View File

@ -55,6 +55,7 @@ import comfy.ldm.pixeldit.model
import comfy.ldm.pixeldit.pid import comfy.ldm.pixeldit.pid
import comfy.ldm.ace.model import comfy.ldm.ace.model
import comfy.ldm.omnigen.omnigen2 import comfy.ldm.omnigen.omnigen2
import comfy.ldm.seedvr.model
import comfy.ldm.boogu.model import comfy.ldm.boogu.model
import comfy.ldm.qwen_image.model import comfy.ldm.qwen_image.model
import comfy.ldm.ideogram4.model import comfy.ldm.ideogram4.model
@ -932,6 +933,17 @@ class HunyuanDiT(BaseModel):
out['image_meta_size'] = comfy.conds.CONDRegular(torch.FloatTensor([[height, width, target_height, target_width, 0, 0]])) out['image_meta_size'] = comfy.conds.CONDRegular(torch.FloatTensor([[height, width, target_height, target_width, 0, 0]]))
return out 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): class PixArt(BaseModel):
def __init__(self, model_config, model_type=ModelType.EPS, device=None): 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) super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.pixart.pixartms.PixArtMS)

View File

@ -470,15 +470,46 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
# PiD (Pixel Diffusion Decoder). Must check BEFORE plain PixelDiT_T2I. # PiD (Pixel Diffusion Decoder). Must check BEFORE plain PixelDiT_T2I.
_lq_w_key = '{}lq_proj.latent_proj.0.weight'.format(key_prefix) _lq_w_key = '{}lq_proj.latent_proj.0.weight'.format(key_prefix)
if _lq_w_key in state_dict_keys: 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) _gate_prefix = '{}lq_proj.gate_modules.'.format(key_prefix)
num_gates = len({k[len(_gate_prefix):].split('.')[0] num_gates = len({k[len(_gate_prefix):].split('.')[0]
for k in state_dict_keys if k.startswith(_gate_prefix)}) 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", dit_config = {"image_model": "pid",
"lq_latent_channels": in_ch, "lq_hidden_dim": hidden_dim}
"latent_spatial_down_factor": 16 if in_ch >= 64 else 8}
if num_gates > 0: if num_gates > 0:
dit_config["lq_interval"] = (14 + num_gates - 1) // num_gates 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 return dit_config
if '{}core.pixel_embedder.proj.weight'.format(key_prefix) in state_dict_keys: # PixelDiT T2I if '{}core.pixel_embedder.proj.weight'.format(key_prefix) in state_dict_keys: # PixelDiT T2I
@ -598,6 +629,44 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
return dit_config 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 if '{}head.modulation'.format(key_prefix) in state_dict_keys: # Wan 2.1
dit_config = {} dit_config = {}
dit_config["image_model"] = "wan2.1" dit_config["image_model"] = "wan2.1"
@ -1119,9 +1188,10 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
return unet_config return unet_config
def model_config_from_unet_config(unet_config, state_dict=None):
def model_config_from_unet_config(unet_config, state_dict=None, unet_key_prefix=""):
for model_config in comfy.supported_models.models: for model_config in comfy.supported_models.models:
if model_config.matches(unet_config, state_dict): if model_config.matches(unet_config, state_dict, unet_key_prefix=unet_key_prefix):
return model_config(unet_config) return model_config(unet_config)
logging.error("no match {}".format(unet_config)) logging.error("no match {}".format(unet_config))
@ -1131,7 +1201,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) unet_config = detect_unet_config(state_dict, unet_key_prefix, metadata=metadata)
if unet_config is None: if unet_config is None:
return None return None
model_config = model_config_from_unet_config(unet_config, state_dict) model_config = model_config_from_unet_config(unet_config, state_dict, unet_key_prefix)
if model_config is None and use_base_if_no_match: if model_config is None and use_base_if_no_match:
model_config = comfy.supported_models_base.BASE(unet_config) model_config = comfy.supported_models_base.BASE(unet_config)

View File

@ -616,6 +616,8 @@ PIN_PRESSURE_HYSTERESIS = 256 * 1024 * 1024
#Freeing registerables on pressure does imply a GPU sync, so go big on #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. #the hysteresis so each expensive sync gives us back a good chunk.
REGISTERABLE_PIN_HYSTERESIS = 2048 * 1024 * 1024 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): def module_size(module):
module_mem = 0 module_mem = 0
@ -642,6 +644,15 @@ def free_pins(size, evict_active=False):
size -= freed size -= freed
return freed_total 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): def ensure_pin_budget(size, evict_active=False):
if args.high_ram: if args.high_ram:
return True return True

View File

@ -1104,6 +1104,21 @@ def _load_quantized_module(module, super_load, state_dict, prefix, local_metadat
scales["convrot_groupsize"] = int( scales["convrot_groupsize"] = int(
layer_conf.get("convrot_groupsize", params_conf.get("convrot_groupsize", 256)) 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: else:
raise ValueError(f"Unsupported quantization format: {module.quant_format}") raise ValueError(f"Unsupported quantization format: {module.quant_format}")
@ -1150,6 +1165,11 @@ def _quantized_weight_state_dict(module, sd, prefix, extra_quant_conf=None, extr
if module.quant_format == "int8_tensorwise" and getattr(params, "convrot", False): if module.quant_format == "int8_tensorwise" and getattr(params, "convrot", False):
quant_conf["convrot"] = True quant_conf["convrot"] = True
quant_conf["convrot_groupsize"] = getattr(params, "convrot_groupsize", 256) 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: if extra_quant_conf:
quant_conf.update(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) sd[f"{prefix}comfy_quant"] = torch.tensor(list(json.dumps(quant_conf).encode("utf-8")), dtype=torch.uint8)
@ -1237,7 +1257,7 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
run_every_op() run_every_op()
input_shape = input.shape input_shape = input.shape
reshaped_3d = False reshaped_nd = False
#If cast needs to apply lora, it should be done in the compute dtype #If cast needs to apply lora, it should be done in the compute dtype
compute_dtype = input.dtype compute_dtype = input.dtype
@ -1274,12 +1294,12 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
# Inference path (unchanged) # Inference path (unchanged)
if _use_quantized and quantize_input: if _use_quantized and quantize_input:
# Reshape 3D tensors to 2D for quantization (needed for NVFP4 and others) # 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 input_reshaped = input.reshape(-1, input_shape[-1]) if input.ndim >= 3 else input
# Fall back to non-quantized for non-2D tensors # Fall back to non-quantized for non-2D tensors
if input_reshaped.ndim == 2: if input_reshaped.ndim == 2:
reshaped_3d = input.ndim == 3 reshaped_nd = input.ndim >= 3
# dtype is now implicit in the layout class # dtype is now implicit in the layout class
scale = getattr(self, 'input_scale', None) scale = getattr(self, 'input_scale', None)
if scale is not None: if scale is not None:
@ -1294,9 +1314,9 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
weight_only_quant=weight_only_quant, weight_only_quant=weight_only_quant,
) )
# Reshape output back to 3D if input was 3D # Reshape output back to original rank if input was >2D
if reshaped_3d: if reshaped_nd:
output = output.reshape((input_shape[0], input_shape[1], self.weight.shape[0])) output = output.reshape((*input_shape[:-1], self.weight.shape[0]))
return output return output
@ -1430,6 +1450,12 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
} }
if hasattr(params, "block_scale"): # NVFP4 if hasattr(params, "block_scale"): # NVFP4
kwargs["block_scale"] = params.block_scale[i] 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)) return QuantizedTensor(weight._qdata[i], weight._layout_cls, type(params)(**kwargs))
def state_dict(self, *args, destination=None, prefix="", **kwargs): def state_dict(self, *args, destination=None, prefix="", **kwargs):

View File

@ -3,6 +3,22 @@ import logging
from comfy.cli_args import args 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: try:
import comfy_kitchen as ck import comfy_kitchen as ck
from comfy_kitchen.tensor import ( from comfy_kitchen.tensor import (
@ -10,6 +26,7 @@ try:
QuantizedLayout, QuantizedLayout,
TensorCoreFP8Layout as _CKFp8Layout, TensorCoreFP8Layout as _CKFp8Layout,
TensorCoreNVFP4Layout as _CKNvfp4Layout, TensorCoreNVFP4Layout as _CKNvfp4Layout,
TensorCoreConvRotW4A4Layout as _CKTensorCoreConvRotW4A4Layout,
TensorWiseINT8Layout as _CKTensorWiseINT8Layout, TensorWiseINT8Layout as _CKTensorWiseINT8Layout,
register_layout_op, register_layout_op,
register_layout_class, register_layout_class,
@ -24,10 +41,22 @@ try:
ck.registry.disable("cuda") ck.registry.disable("cuda")
logging.warning("WARNING: You need pytorch with cu130 or higher to use optimized CUDA operations.") logging.warning("WARNING: You need pytorch with cu130 or higher to use optimized CUDA operations.")
if args.enable_triton_backend: # 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()):
try: try:
import triton import triton
logging.info("Found triton %s. Enabling comfy-kitchen triton backend.", triton.__version__) 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")
except ImportError as e: except ImportError as e:
logging.error(f"Failed to import triton, Error: {e}, the comfy-kitchen triton backend will not be available.") logging.error(f"Failed to import triton, Error: {e}, the comfy-kitchen triton backend will not be available.")
ck.registry.disable("triton") ck.registry.disable("triton")
@ -51,6 +80,9 @@ except ImportError as e:
class _CKTensorWiseINT8Layout: class _CKTensorWiseINT8Layout:
pass pass
class _CKTensorCoreConvRotW4A4Layout:
pass
def register_layout_class(name, cls): def register_layout_class(name, cls):
pass pass
@ -179,6 +211,7 @@ class TensorCoreFP8E5M2Layout(_TensorCoreFP8LayoutBase):
# Backward compatibility alias - default to E4M3 # Backward compatibility alias - default to E4M3
TensorCoreFP8Layout = TensorCoreFP8E4M3Layout TensorCoreFP8Layout = TensorCoreFP8E4M3Layout
TensorWiseINT8Layout = _CKTensorWiseINT8Layout TensorWiseINT8Layout = _CKTensorWiseINT8Layout
TensorCoreConvRotW4A4Layout = _CKTensorCoreConvRotW4A4Layout
# ============================================================================== # ==============================================================================
@ -190,6 +223,7 @@ register_layout_class("TensorCoreFP8E4M3Layout", TensorCoreFP8E4M3Layout)
register_layout_class("TensorCoreFP8E5M2Layout", TensorCoreFP8E5M2Layout) register_layout_class("TensorCoreFP8E5M2Layout", TensorCoreFP8E5M2Layout)
register_layout_class("TensorCoreNVFP4Layout", TensorCoreNVFP4Layout) register_layout_class("TensorCoreNVFP4Layout", TensorCoreNVFP4Layout)
register_layout_class("TensorWiseINT8Layout", _CKTensorWiseINT8Layout) register_layout_class("TensorWiseINT8Layout", _CKTensorWiseINT8Layout)
register_layout_class("TensorCoreConvRotW4A4Layout", _CKTensorCoreConvRotW4A4Layout)
if _CK_MXFP8_AVAILABLE: if _CK_MXFP8_AVAILABLE:
register_layout_class("TensorCoreMXFP8Layout", TensorCoreMXFP8Layout) register_layout_class("TensorCoreMXFP8Layout", TensorCoreMXFP8Layout)
@ -227,6 +261,13 @@ QUANT_ALGOS["int8_tensorwise"] = {
"quantize_input": False, "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 # Re-exports for backward compatibility
@ -239,6 +280,7 @@ __all__ = [
"TensorCoreFP8E4M3Layout", "TensorCoreFP8E4M3Layout",
"TensorCoreFP8E5M2Layout", "TensorCoreFP8E5M2Layout",
"TensorCoreNVFP4Layout", "TensorCoreNVFP4Layout",
"TensorCoreConvRotW4A4Layout",
"TensorWiseINT8Layout", "TensorWiseINT8Layout",
"QUANT_ALGOS", "QUANT_ALGOS",
"register_layout_op", "register_layout_op",

View File

@ -16,6 +16,7 @@ import comfy.ldm.cosmos.vae
import comfy.ldm.wan.vae import comfy.ldm.wan.vae
import comfy.ldm.wan.vae2_2 import comfy.ldm.wan.vae2_2
import comfy.ldm.hunyuan3d.vae import comfy.ldm.hunyuan3d.vae
import comfy.ldm.seedvr.vae
import comfy.ldm.triposplat.vae import comfy.ldm.triposplat.vae
import comfy.ldm.ace.vae.music_dcae_pipeline import comfy.ldm.ace.vae.music_dcae_pipeline
import comfy.ldm.cogvideo.vae import comfy.ldm.cogvideo.vae
@ -473,7 +474,8 @@ class CLIP:
class VAE: class VAE:
def __init__(self, sd=None, device=None, config=None, dtype=None, metadata=None): def __init__(self, sd=None, device=None, config=None, dtype=None, metadata=None):
if 'decoder.up_blocks.0.resnets.0.norm1.weight' in sd.keys(): #diffusers format 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
sd = diffusers_convert.convert_vae_state_dict(sd) sd = diffusers_convert.convert_vae_state_dict(sd)
if model_management.is_amd(): if model_management.is_amd():
@ -500,6 +502,8 @@ class VAE:
self.upscale_index_formula = None self.upscale_index_formula = None
self.extra_1d_channel = None self.extra_1d_channel = None
self.crop_input = True self.crop_input = True
self.handles_tiling = False
self.format_encoded = None
self.audio_sample_rate = 44100 self.audio_sample_rate = 44100
@ -546,6 +550,22 @@ class VAE:
self.first_stage_model = StageC_coder() self.first_stage_model = StageC_coder()
self.downscale_ratio = 32 self.downscale_ratio = 32
self.latent_channels = 16 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: elif "decoder.conv_in.weight" in sd:
if sd['decoder.conv_in.weight'].shape[1] == 64: 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} 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}
@ -1012,6 +1032,10 @@ 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()) 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)) 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): 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, 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) 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)
@ -1048,6 +1072,25 @@ 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()) 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) 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={}): def decode(self, samples_in, vae_options={}):
self.throw_exception_if_invalid() self.throw_exception_if_invalid()
pixel_samples = None pixel_samples = None
@ -1095,11 +1138,19 @@ class VAE:
if dims == 1 or self.extra_1d_channel is not None: if dims == 1 or self.extra_1d_channel is not None:
pixel_samples = self.decode_tiled_1d(samples_in) pixel_samples = self.decode_tiled_1d(samples_in)
elif dims == 2: elif dims == 2:
pixel_samples = self.decode_tiled_(samples_in) 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)
elif dims == 3: elif dims == 3:
tile = 256 // self.spacial_compression_decode() tile = 256 // self.spacial_compression_decode()
overlap = tile // 4 overlap = tile // 4
pixel_samples = self.decode_tiled_3d(samples_in, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap)) 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 = pixel_samples.to(self.output_device).movedim(1,-1) pixel_samples = pixel_samples.to(self.output_device).movedim(1,-1)
return pixel_samples return pixel_samples
@ -1118,7 +1169,9 @@ class VAE:
args["overlap"] = overlap args["overlap"] = overlap
with model_management.cuda_device_context(self.device): with model_management.cuda_device_context(self.device):
if dims == 1 or self.extra_1d_channel is not None: 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:
args.pop("tile_y") args.pop("tile_y")
output = self.decode_tiled_1d(samples, **args) output = self.decode_tiled_1d(samples, **args)
elif dims == 2: elif dims == 2:
@ -1179,12 +1232,17 @@ class VAE:
if self.latent_dim == 3: if self.latent_dim == 3:
tile = 256 tile = 256
overlap = tile // 4 overlap = tile // 4
samples = self.encode_tiled_3d(pixel_samples, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap)) 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))
elif self.latent_dim == 1 or self.extra_1d_channel is not None: elif self.latent_dim == 1 or self.extra_1d_channel is not None:
samples = self.encode_tiled_1d(pixel_samples) samples = self.encode_tiled_1d(pixel_samples)
else: else:
samples = self.encode_tiled_(pixel_samples) samples = self.encode_tiled_(pixel_samples)
if self.format_encoded is not None:
samples = self.format_encoded(samples)
return samples return samples
def encode_tiled(self, pixel_samples, tile_x=None, tile_y=None, overlap=None, tile_t=None, overlap_t=None): def encode_tiled(self, pixel_samples, tile_x=None, tile_y=None, overlap=None, tile_t=None, overlap_t=None):
@ -1192,7 +1250,7 @@ class VAE:
pixel_samples = self.vae_encode_crop_pixels(pixel_samples) pixel_samples = self.vae_encode_crop_pixels(pixel_samples)
dims = self.latent_dim dims = self.latent_dim
pixel_samples = pixel_samples.movedim(-1, 1) pixel_samples = pixel_samples.movedim(-1, 1)
if dims == 3: if dims == 3 and pixel_samples.ndim < 5:
if not self.not_video: if not self.not_video:
pixel_samples = pixel_samples.movedim(1, 0).unsqueeze(0) pixel_samples = pixel_samples.movedim(1, 0).unsqueeze(0)
else: else:
@ -1216,21 +1274,27 @@ class VAE:
elif dims == 2: elif dims == 2:
samples = self.encode_tiled_(pixel_samples, **args) samples = self.encode_tiled_(pixel_samples, **args)
elif dims == 3: elif dims == 3:
if tile_t is not None: if self.handles_tiling:
tile_t_latent = max(2, self.downscale_ratio[0](tile_t)) samples = self._encode_tiled_owned(pixel_samples, **self._owned_tiled_args(tile_x, tile_y, overlap, tile_t, overlap_t))
else: else:
tile_t_latent = 9999 if tile_t is not None:
args["tile_t"] = self.upscale_ratio[0](tile_t_latent) 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)
if overlap_t is None: spatial_overlap = overlap if overlap is not None else 64
args["overlap"] = (1, overlap, overlap) if overlap_t is None:
else: args["overlap"] = (1, spatial_overlap, spatial_overlap)
args["overlap"] = (self.upscale_ratio[0](max(1, min(tile_t_latent // 2, self.downscale_ratio[0](overlap_t)))), overlap, overlap) else:
maximum = pixel_samples.shape[2] args["overlap"] = (self.upscale_ratio[0](max(1, min(tile_t_latent // 2, self.downscale_ratio[0](overlap_t)))), spatial_overlap, spatial_overlap)
maximum = self.upscale_ratio[0](self.downscale_ratio[0](maximum)) 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 return samples
def get_sd(self): def get_sd(self):
@ -1898,7 +1962,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) manual_cast_dtype = model_management.unet_manual_cast(None, load_device, model_config.supported_inference_dtypes)
else: else:
manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes) 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) model_config.set_inference_dtype(unet_dtype, manual_cast_dtype, device=load_device)
if model_config.clip_vision_prefix is not None: if model_config.clip_vision_prefix is not None:
if output_clipvision: if output_clipvision:
@ -2039,7 +2103,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) manual_cast_dtype = model_management.unet_manual_cast(None, load_device, model_config.supported_inference_dtypes)
else: else:
manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes) 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) model_config.set_inference_dtype(unet_dtype, manual_cast_dtype, device=load_device)
if custom_operations is not None: if custom_operations is not None:
model_config.custom_operations = custom_operations model_config.custom_operations = custom_operations

View File

@ -1685,6 +1685,40 @@ class Chroma(supported_models_base.BASE):
t5_detect = comfy.text_encoders.sd3_clip.t5_xxl_detect(state_dict, "{}t5xxl.transformer.".format(pref)) 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)) 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): class ChromaRadiance(Chroma):
unet_config = { unet_config = {
"image_model": "chroma_radiance", "image_model": "chroma_radiance",
@ -2348,6 +2382,7 @@ models = [
HiDream, HiDream,
HiDreamO1, HiDreamO1,
Chroma, Chroma,
SeedVR2,
ChromaRadiance, ChromaRadiance,
ACEStep, ACEStep,
ACEStep15, ACEStep15,

View File

@ -54,13 +54,13 @@ class BASE:
optimizations = {"fp8": False} optimizations = {"fp8": False}
@classmethod @classmethod
def matches(s, unet_config, state_dict=None): def matches(s, unet_config, state_dict=None, unet_key_prefix=""):
for k in s.unet_config: for k in s.unet_config:
if k not in unet_config or s.unet_config[k] != unet_config[k]: if k not in unet_config or s.unet_config[k] != unet_config[k]:
return False return False
if state_dict is not None: if state_dict is not None:
for k in s.required_keys: for k in s.required_keys:
if k not in state_dict: if k.format(unet_key_prefix) not in state_dict:
return False return False
return True return True
@ -115,7 +115,7 @@ class BASE:
replace_prefix = {"": self.vae_key_prefix[0]} replace_prefix = {"": self.vae_key_prefix[0]}
return utils.state_dict_prefix_replace(state_dict, replace_prefix) return utils.state_dict_prefix_replace(state_dict, replace_prefix)
def set_inference_dtype(self, dtype, manual_cast_dtype): def set_inference_dtype(self, dtype, manual_cast_dtype, device=None):
self.unet_config['dtype'] = dtype self.unet_config['dtype'] = dtype
self.manual_cast_dtype = manual_cast_dtype self.manual_cast_dtype = manual_cast_dtype

View File

@ -1088,7 +1088,7 @@ class Gemma4_Tokenizer():
h, w = samples.shape[2], samples.shape[3] h, w = samples.shape[2], samples.shape[3]
patch_size = 16 patch_size = 16
pooling_k = 3 pooling_k = 3
max_soft_tokens = 70 if is_video else 280 # video uses smaller token budget per frame max_soft_tokens = kwargs.get("max_soft_tokens", 70 if is_video else 280)
max_patches = max_soft_tokens * pooling_k * pooling_k max_patches = max_soft_tokens * pooling_k * pooling_k
target_px = max_patches * patch_size * patch_size target_px = max_patches * patch_size * patch_size
factor = (target_px / (h * w)) ** 0.5 factor = (target_px / (h * w)) ** 0.5

View File

@ -17,6 +17,10 @@ class Seedream4Options(BaseModel):
max_images: int = Field(15) max_images: int = Field(15)
class Seedream5OptimizePromptOptions(BaseModel):
thinking: Literal["auto", "enabled", "disabled"] = Field(...)
class Seedream4TaskCreationRequest(BaseModel): class Seedream4TaskCreationRequest(BaseModel):
model: str = Field(...) model: str = Field(...)
prompt: str = Field(...) prompt: str = Field(...)
@ -28,6 +32,7 @@ class Seedream4TaskCreationRequest(BaseModel):
sequential_image_generation_options: Seedream4Options | None = Field(Seedream4Options(max_images=15)) sequential_image_generation_options: Seedream4Options | None = Field(Seedream4Options(max_images=15))
watermark: bool = Field(False) watermark: bool = Field(False)
output_format: str | None = None output_format: str | None = None
optimize_prompt_options: Seedream5OptimizePromptOptions | None = None
class ImageTaskCreationResponse(BaseModel): class ImageTaskCreationResponse(BaseModel):

View File

@ -77,6 +77,7 @@ class To3DUVTaskRequest(BaseModel):
class To3DPartTaskRequest(BaseModel): class To3DPartTaskRequest(BaseModel):
File: TaskFile3DInput = Field(...) File: TaskFile3DInput = Field(...)
EnableStagedGeneration: bool | None = Field(None)
class TextureEditImageInfo(BaseModel): class TextureEditImageInfo(BaseModel):

View File

@ -0,0 +1,49 @@
from pydantic import BaseModel, Field
class SyncInputItem(BaseModel):
type: str = Field(..., description="Input kind: 'video', 'image' or 'audio'.")
url: str = Field(...)
class SyncActiveSpeakerDetection(BaseModel):
auto_detect: bool | None = Field(
None, description="Detect the active speaker automatically. Video input only; rejected for images."
)
frame_number: int | None = Field(
None, description="Frame used for manual speaker selection. Must be 0 for image inputs."
)
coordinates: list[int] | None = Field(
None, description="Pixel [x, y] of the speaker's face in the frame selected by frame_number."
)
class SyncGenerationOptions(BaseModel):
sync_mode: str | None = Field(
None,
description="How to resolve an audio/video duration mismatch: "
"cut_off, bounce, loop, silence or remap. Ignored for image inputs.",
)
i2v_prompt: str | None = Field(
None, description="Motion prompt for image-to-video generation. Image input only."
)
active_speaker_detection: SyncActiveSpeakerDetection | None = Field(None)
class SyncGenerationRequest(BaseModel):
model: str = Field(..., description="Generation model, e.g. 'sync-3'.")
input: list[SyncInputItem] = Field(
..., description="Exactly one visual input (video or image) plus one audio input."
)
options: SyncGenerationOptions | None = Field(None)
class SyncGeneration(BaseModel):
"""Subset of the Generation object returned by POST /v2/generate and GET /v2/generate/{id}."""
id: str = Field(...)
status: str = Field(..., description="PENDING | PROCESSING | COMPLETED | FAILED | REJECTED")
outputUrl: str | None = Field(None)
outputDuration: float | None = Field(None)
error: str | None = Field(None, description="Human-readable failure message.")
errorCode: str | None = Field(None, description="Stable machine-readable code from the GET /v2/errors catalog.")

View File

@ -34,6 +34,7 @@ from comfy_api_nodes.apis.bytedance import (
SeedanceVirtualLibraryCreateAssetRequest, SeedanceVirtualLibraryCreateAssetRequest,
Seedream4Options, Seedream4Options,
Seedream4TaskCreationRequest, Seedream4TaskCreationRequest,
Seedream5OptimizePromptOptions,
TaskAudioContent, TaskAudioContent,
TaskAudioContentUrl, TaskAudioContentUrl,
TaskCreationResponse, TaskCreationResponse,
@ -875,6 +876,17 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode):
tooltip='Whether to add an "AI generated" watermark to the image.', tooltip='Whether to add an "AI generated" watermark to the image.',
advanced=True, 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=[ outputs=[
IO.Image.Output(), IO.Image.Output(),
@ -920,6 +932,7 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode):
model: dict, model: dict,
seed: int = 0, seed: int = 0,
watermark: bool = False, watermark: bool = False,
thinking: bool = True,
) -> IO.NodeOutput: ) -> IO.NodeOutput:
validate_string(prompt, strip_whitespace=True, min_length=1) validate_string(prompt, strip_whitespace=True, min_length=1)
model_id = SEEDREAM_MODELS[model["model"]] model_id = SEEDREAM_MODELS[model["model"]]
@ -979,6 +992,10 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode):
raise ValueError( raise ValueError(
"The maximum number of generated images plus the number of reference images cannot exceed 15." "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] = [] reference_images_urls: list[str] = []
if image_tensors: if image_tensors:
@ -992,6 +1009,9 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode):
wait_label="Uploading reference images", 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( response = await sync_op(
cls, cls,
ApiEndpoint(path=BYTEPLUS_IMAGE_ENDPOINT, method="POST"), ApiEndpoint(path=BYTEPLUS_IMAGE_ENDPOINT, method="POST"),
@ -1005,6 +1025,7 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode):
sequential_image_generation=None if is_pro else sequential_image_generation, 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), sequential_image_generation_options=None if is_pro else Seedream4Options(max_images=max_images),
watermark=watermark, watermark=watermark,
optimize_prompt_options=optimize_prompt_options,
), ),
) )
if len(response.data) == 1: if len(response.data) == 1:

View File

@ -1133,7 +1133,9 @@ class GeminiImage2(IO.ComfyNode):
) -> IO.NodeOutput: ) -> IO.NodeOutput:
validate_string(prompt, strip_whitespace=True, min_length=1) validate_string(prompt, strip_whitespace=True, min_length=1)
if model == "Nano Banana 2 (Gemini 3.1 Flash Image)": 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)] parts: list[GeminiPart] = [GeminiPart(text=prompt)]
if images is not None: if images is not None:
@ -1507,7 +1509,7 @@ class GeminiNanoBanana2V2(IO.ComfyNode):
validate_string(prompt, strip_whitespace=True, min_length=1) validate_string(prompt, strip_whitespace=True, min_length=1)
model_choice = model["model"] model_choice = model["model"]
if model_choice == "Nano Banana 2 (Gemini 3.1 Flash Image)": 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": elif model_choice == "Nano Banana 2 Lite":
model_id = "gemini-3.1-flash-lite-image" model_id = "gemini-3.1-flash-lite-image"
else: else:

View File

@ -642,6 +642,7 @@ class Tencent3DPartNode(IO.ComfyNode):
response_model=To3DProTaskCreateResponse, response_model=To3DProTaskCreateResponse,
data=To3DPartTaskRequest( data=To3DPartTaskRequest(
File=TaskFile3DInput(Type=file_format.upper(), Url=model_url), File=TaskFile3DInput(Type=file_format.upper(), Url=model_url),
EnableStagedGeneration=True,
), ),
is_rate_limited=_is_tencent_rate_limited, is_rate_limited=_is_tencent_rate_limited,
) )

View File

@ -0,0 +1,391 @@
from typing_extensions import override
from comfy_api.latest import IO, ComfyExtension, Input
from comfy_api_nodes.apis.sync_so import (
SyncActiveSpeakerDetection,
SyncGeneration,
SyncGenerationOptions,
SyncGenerationRequest,
SyncInputItem,
)
from comfy_api_nodes.util import (
ApiEndpoint,
download_url_to_video_output,
downscale_image_tensor,
downscale_image_tensor_by_max_side,
get_image_dimensions,
get_number_of_images,
poll_op,
sync_op,
upload_audio_to_comfyapi,
upload_image_to_comfyapi,
upload_video_to_comfyapi,
validate_audio_duration,
)
class SyncLipSyncNode(IO.ComfyNode):
@classmethod
def define_schema(cls) -> IO.Schema:
return IO.Schema(
node_id="SyncLipSyncNode",
display_name="sync.so Lip Sync",
category="partner/video/sync.so",
description=(
"Re-sync mouth movement in a video to new speech audio using sync.so. "
"Handles close-ups, profiles and obstructions automatically while preserving "
"the speaker's expression. Cost scales with output duration."
),
inputs=[
IO.Video.Input(
"video",
tooltip="Footage of the speaker to re-sync. Up to 4K (4096x2160); "
"a constant frame rate of 24/25/30 fps works best.",
),
IO.Audio.Input(
"audio",
tooltip="Speech audio to sync the mouth to.",
),
IO.Int.Input(
"seed",
default=42,
min=0,
max=2147483647,
control_after_generate=True,
tooltip="Seed controls whether the node should re-run; "
"results are non-deterministic regardless of seed.",
),
IO.DynamicCombo.Input(
"model",
options=[
IO.DynamicCombo.Option(
"sync-3",
[
IO.Combo.Input(
"sync_mode",
options=["bounce", "cut_off", "loop", "silence", "remap"],
default="bounce",
tooltip=(
"How to handle a duration mismatch between video and audio; "
"this also sets the output length. "
"bounce: video plays forward then backward until the audio ends "
"(output = audio length). "
"loop: video restarts until the audio ends (output = audio length). "
"remap: video is time-stretched to match the audio (output = audio length). "
"cut_off: the longer track is trimmed (output = shorter length). "
"silence: nothing is trimmed; the shorter track is padded "
"(output = longer length)."
),
),
IO.Combo.Input(
"speaker_selection",
options=["default", "auto-detect", "coordinates"],
default="default",
tooltip=(
"Which face to lipsync when several people are visible. "
"default: let the model decide. "
"auto-detect: detect and follow the active speaker. "
"coordinates: target the face at pixel (speaker_x, speaker_y) "
"in the frame chosen by speaker_frame."
),
),
IO.Int.Input(
"speaker_frame",
default=0,
min=0,
max=1_000_000,
advanced=True,
tooltip="Video frame used to locate the speaker. "
"Only used when speaker_selection is 'coordinates'.",
),
IO.Int.Input(
"speaker_x",
default=0,
min=0,
max=4096,
advanced=True,
tooltip="X pixel coordinate of the speaker's face. "
"Only used when speaker_selection is 'coordinates'.",
),
IO.Int.Input(
"speaker_y",
default=0,
min=0,
max=4096,
advanced=True,
tooltip="Y pixel coordinate of the speaker's face. "
"Only used when speaker_selection is 'coordinates'.",
),
],
)
],
tooltip="sync.so generation model.",
),
],
outputs=[IO.Video.Output()],
hidden=[
IO.Hidden.auth_token_comfy_org,
IO.Hidden.api_key_comfy_org,
IO.Hidden.unique_id,
],
is_api_node=True,
price_badge=IO.PriceBadge(
expr="""{"type":"usd","usd":0.19019,"format":{"approximate":true,"suffix":"/second"}}""",
),
)
@classmethod
async def execute(
cls,
video: Input.Video,
audio: Input.Audio,
seed: int,
model: dict,
) -> IO.NodeOutput:
try:
width, height = video.get_dimensions()
except Exception:
width = height = None
if width and height and (max(width, height) > 4096 or width * height > 4096 * 2160):
raise ValueError(
f"sync.so rejects videos above 4K (4096x2160); got {width}x{height}. Downscale the video first."
)
validate_audio_duration(audio, max_duration=600)
if model["speaker_selection"] == "auto-detect":
speaker_detection = SyncActiveSpeakerDetection(auto_detect=True)
elif model["speaker_selection"] == "coordinates":
speaker_detection = SyncActiveSpeakerDetection(
frame_number=model["speaker_frame"],
coordinates=[model["speaker_x"], model["speaker_y"]],
)
else:
speaker_detection = None
video_url = await upload_video_to_comfyapi(cls, video, max_duration=600)
audio_url = await upload_audio_to_comfyapi(cls, audio)
generation = await sync_op(
cls,
ApiEndpoint(path="/proxy/synclabs/v2/generate", method="POST"),
response_model=SyncGeneration,
data=SyncGenerationRequest(
model=model["model"],
input=[
SyncInputItem(type="video", url=video_url),
SyncInputItem(type="audio", url=audio_url),
],
options=SyncGenerationOptions(
sync_mode=model["sync_mode"],
active_speaker_detection=speaker_detection,
),
),
)
generation = await poll_op(
cls,
ApiEndpoint(path=f"/proxy/synclabs/v2/generate/{generation.id}"),
response_model=SyncGeneration,
status_extractor=lambda g: g.status,
completed_statuses=["COMPLETED", "FAILED", "REJECTED"],
failed_statuses=[],
queued_statuses=["PENDING"],
poll_interval=10.0,
)
if generation.status != "COMPLETED":
code = f" [{generation.errorCode}]" if generation.errorCode else ""
raise ValueError(
f"sync.so generation {generation.status.lower()}{code}: "
f"{generation.error or 'no error details provided'}"
)
if not generation.outputUrl:
raise ValueError("sync.so generation completed but no output URL was returned.")
return IO.NodeOutput(await download_url_to_video_output(generation.outputUrl))
class SyncTalkingImageNode(IO.ComfyNode):
@classmethod
def define_schema(cls) -> IO.Schema:
return IO.Schema(
node_id="SyncTalkingImageNode",
display_name="sync.so Talking Image",
category="partner/video/sync.so",
description=(
"Animate a still portrait into a talking video driven by speech audio, "
"using sync.so's sync-3 model. The output duration matches the audio. "
"Cost scales with output duration."
),
inputs=[
IO.Image.Input(
"image",
tooltip="A single image with a clearly visible face, up to 4K (4096x2160).",
),
IO.Audio.Input(
"audio",
tooltip="Speech audio driving the talking video; the output duration matches it. "
"Chain any TTS node here to drive the animation from text.",
),
IO.String.Input(
"prompt",
multiline=True,
default="",
tooltip="Optional guidance for how the portrait comes to life, e.g. "
"'make the subject smile and look at the camera'. "
"Leave empty for natural talking motion.",
),
IO.Int.Input(
"seed",
default=0,
min=0,
max=2147483647,
control_after_generate=True,
tooltip="Seed controls whether the node should re-run; "
"results are non-deterministic regardless of seed.",
),
IO.DynamicCombo.Input(
"model",
options=[
IO.DynamicCombo.Option(
"sync-3",
[
IO.Combo.Input(
"speaker_selection",
options=["default", "coordinates"],
default="default",
tooltip=(
"Which face to animate when several people are visible. "
"default: let the model decide. "
"coordinates: target the face at pixel (speaker_x, speaker_y) "
"in the image. Auto-detection is not supported for images."
),
),
IO.Int.Input(
"speaker_x",
default=0,
min=0,
max=4096,
advanced=True,
tooltip="X pixel coordinate of the speaker's face. "
"Only used when speaker_selection is 'coordinates'.",
),
IO.Int.Input(
"speaker_y",
default=0,
min=0,
max=4096,
advanced=True,
tooltip="Y pixel coordinate of the speaker's face. "
"Only used when speaker_selection is 'coordinates'.",
),
IO.Boolean.Input(
"auto_downscale",
default=True,
advanced=True,
tooltip="Automatically downscale the image if it exceeds the 4K "
"(4096x2160) input limit; speaker coordinates are scaled to match. "
"When disabled, an oversized image raises an error instead.",
),
],
)
],
tooltip="sync.so generation model. Image input is exclusive to sync-3.",
),
],
outputs=[IO.Video.Output()],
hidden=[
IO.Hidden.auth_token_comfy_org,
IO.Hidden.api_key_comfy_org,
IO.Hidden.unique_id,
],
is_api_node=True,
price_badge=IO.PriceBadge(
expr="""{"type":"usd","usd":0.19019,"format":{"approximate":true,"suffix":"/second"}}""",
),
)
@classmethod
async def execute(
cls,
image: Input.Image,
audio: Input.Audio,
prompt: str,
seed: int,
model: dict,
) -> IO.NodeOutput:
if get_number_of_images(image) != 1:
raise ValueError("Exactly one image is required; got a batch. Pick one frame first.")
validate_audio_duration(audio, max_duration=600)
height, width = get_image_dimensions(image)
speaker_x, speaker_y = model["speaker_x"], model["speaker_y"]
if max(width, height) > 4096 or width * height > 4096 * 2160:
if not model["auto_downscale"]:
raise ValueError(
f"sync.so rejects images above 4K (4096x2160); got {width}x{height}. "
"Downscale the image first or enable auto_downscale."
)
image = downscale_image_tensor(image, total_pixels=4096 * 2160)
image = downscale_image_tensor_by_max_side(image, max_side=4096)
new_height, new_width = get_image_dimensions(image)
# speaker coordinates are given in the original image's pixel space
speaker_x = min(new_width - 1, round(speaker_x * new_width / width))
speaker_y = min(new_height - 1, round(speaker_y * new_height / height))
if model["speaker_selection"] == "coordinates":
speaker_detection = SyncActiveSpeakerDetection(
frame_number=0, # images have a single frame; auto_detect is rejected by the API
coordinates=[speaker_x, speaker_y],
)
else:
speaker_detection = None
image_url = await upload_image_to_comfyapi(cls, image, mime_type="image/png", total_pixels=None)
audio_url = await upload_audio_to_comfyapi(cls, audio)
generation = await sync_op(
cls,
ApiEndpoint(path="/proxy/synclabs/v2/generate", method="POST"),
response_model=SyncGeneration,
data=SyncGenerationRequest(
model=model["model"],
input=[
SyncInputItem(type="image", url=image_url),
SyncInputItem(type="audio", url=audio_url),
],
options=SyncGenerationOptions(
i2v_prompt=prompt.strip() or None,
active_speaker_detection=speaker_detection,
),
),
)
generation = await poll_op(
cls,
ApiEndpoint(path=f"/proxy/synclabs/v2/generate/{generation.id}"),
response_model=SyncGeneration,
status_extractor=lambda g: g.status,
completed_statuses=["COMPLETED", "FAILED", "REJECTED"],
failed_statuses=[],
queued_statuses=["PENDING"],
poll_interval=10.0,
)
if generation.status != "COMPLETED":
code = f" [{generation.errorCode}]" if generation.errorCode else ""
raise ValueError(
f"sync.so generation {generation.status.lower()}{code}: "
f"{generation.error or 'no error details provided'}"
)
if not generation.outputUrl:
raise ValueError("sync.so generation completed but no output URL was returned.")
return IO.NodeOutput(await download_url_to_video_output(generation.outputUrl))
class SyncExtension(ComfyExtension):
@override
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
return [
SyncLipSyncNode,
SyncTalkingImageNode,
]
async def comfy_entrypoint() -> SyncExtension:
return SyncExtension()

View File

@ -15,6 +15,7 @@ from comfy.comfy_api_env import normalize_comfy_api_base
from comfy.deploy_environment import get_deploy_environment from comfy.deploy_environment import get_deploy_environment
from comfy.model_management import processing_interrupted from comfy.model_management import processing_interrupted
from comfy_api.latest import IO from comfy_api.latest import IO
from comfyui_version import __version__ as comfyui_version
from .common_exceptions import ProcessingInterrupted 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), **get_auth_header(node_cls),
"Comfy-Env": get_deploy_environment(), "Comfy-Env": get_deploy_environment(),
"Comfy-Usage-Source": get_usage_source(node_cls), "Comfy-Usage-Source": get_usage_source(node_cls),
"Comfy-Core-Version": comfyui_version,
} }

View File

@ -503,6 +503,8 @@ RAM_CACHE_DEFAULT_RAM_USAGE = 0.05
RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER = 1.3 RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER = 1.3
RAM_CACHE_LARGE_INTERMEDIATE = 512 * 1024 ** 2
def all_outputs_dynamic(outputs): def all_outputs_dynamic(outputs):
if outputs is None: if outputs is None:
@ -517,7 +519,6 @@ def all_outputs_dynamic(outputs):
return True return True
class RAMPressureCache(LRUCache): class RAMPressureCache(LRUCache):
def __init__(self, key_class, enable_providers=False): def __init__(self, key_class, enable_providers=False):
@ -539,9 +540,9 @@ class RAMPressureCache(LRUCache):
self.timestamps[self.cache_key_set.get_data_key(node_id)] = time.time() self.timestamps[self.cache_key_set.get_data_key(node_id)] = time.time()
super().set_local(node_id, value) super().set_local(node_id, value)
def ram_release(self, target, free_active=False): def ram_release(self, target, free_active=False, min_entry_size=0):
if psutil.virtual_memory().available >= target: if psutil.virtual_memory().available >= target:
return return 0
clean_list = [] clean_list = []
@ -555,8 +556,9 @@ class RAMPressureCache(LRUCache):
oom_score = RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER ** (self.generation - self.used_generation[key]) oom_score = RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER ** (self.generation - self.used_generation[key])
ram_usage = RAM_CACHE_DEFAULT_RAM_USAGE ram_usage = RAM_CACHE_DEFAULT_RAM_USAGE
oom_ram_usage = ram_usage
def scan_list_for_ram_usage(outputs): def scan_list_for_ram_usage(outputs):
nonlocal ram_usage nonlocal ram_usage, oom_ram_usage
if outputs is None: if outputs is None:
return return
for output in outputs: for output in outputs:
@ -564,19 +566,26 @@ class RAMPressureCache(LRUCache):
scan_list_for_ram_usage(output) scan_list_for_ram_usage(output)
elif isinstance(output, torch.Tensor) and output.device.type == 'cpu': elif isinstance(output, torch.Tensor) and output.device.type == 'cpu':
ram_usage += output.numel() * output.element_size() 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: elif isinstance(output, ModelPatcher) and self.used_generation[key] != self.generation:
#old ModelPatchers are the first to go #old ModelPatchers are the first to go
ram_usage = 1e30 oom_ram_usage = 1e30
scan_list_for_ram_usage(cache_entry.outputs) scan_list_for_ram_usage(cache_entry.outputs)
oom_score *= ram_usage if ram_usage < min_entry_size:
continue
oom_score *= oom_ram_usage
#In the case where we have no information on the node ram usage at all, #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) #break OOM score ties on the last touch timestamp (pure LRU)
bisect.insort(clean_list, (oom_score, self.timestamps[key], key)) bisect.insort(clean_list, (oom_score, self.timestamps[key], key, ram_usage))
freed = 0
while psutil.virtual_memory().available < target and clean_list: while psutil.virtual_memory().available < target and clean_list:
_, _, key = clean_list.pop() _, _, key, ram_usage = clean_list.pop()
del self.cache[key] del self.cache[key]
self.used_generation.pop(key, None) self.used_generation.pop(key, None)
self.timestamps.pop(key, None) self.timestamps.pop(key, None)
self.children.pop(key, None) self.children.pop(key, None)
freed += ram_usage
return freed

View File

@ -56,6 +56,9 @@ PREVIEWABLE_MEDIA_TYPES = frozenset({'images', 'video', 'audio', '3d', 'text'})
# 3D file extensions for preview fallback (no dedicated media_type exists) # 3D file extensions for preview fallback (no dedicated media_type exists)
THREE_D_EXTENSIONS = frozenset({'.obj', '.fbx', '.gltf', '.glb', '.usdz'}) 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: def has_3d_extension(filename: str) -> bool:
lower = filename.lower() lower = filename.lower()
@ -143,9 +146,10 @@ def is_previewable(media_type: str, item: dict) -> bool:
Maintains backwards compatibility with existing logic. Maintains backwards compatibility with existing logic.
Priority: Priority:
1. media_type is 'images', 'video', 'audio', or '3d' 1. media_type is 'images', 'video', 'audio', '3d', or 'text'
2. format field starts with 'video/' or 'audio/' 2. format field starts with 'video/' or 'audio/'
3. filename has a 3D extension (.obj, .fbx, .gltf, .glb, .usdz) 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: if media_type in PREVIEWABLE_MEDIA_TYPES:
return True return True
@ -156,10 +160,12 @@ def is_previewable(media_type: str, item: dict) -> bool:
if fmt and (fmt.startswith('video/') or fmt.startswith('audio/')): if fmt and (fmt.startswith('video/') or fmt.startswith('audio/')):
return True return True
# Check for 3D files by extension # Check for 3D and text files by extension
filename = item.get('filename', '').lower() filename = item.get('filename', '').lower()
if any(filename.endswith(ext) for ext in THREE_D_EXTENSIONS): if any(filename.endswith(ext) for ext in THREE_D_EXTENSIONS):
return True return True
if any(filename.endswith(ext) for ext in TEXT_EXTENSIONS):
return True
return False return False
@ -255,6 +261,10 @@ def get_outputs_summary(outputs: dict) -> tuple[int, Optional[dict]]:
Preview priority (matching frontend): Preview priority (matching frontend):
1. type="output" with previewable media 1. type="output" with previewable media
2. Any 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 count = 0
preview_output = None preview_output = None
@ -275,7 +285,6 @@ def get_outputs_summary(outputs: dict) -> tuple[int, Optional[dict]]:
if normalized is None: if normalized is None:
# Not a 3D file string — check for text preview # Not a 3D file string — check for text preview
if media_type == 'text': if media_type == 'text':
count += 1
if preview_output is None: if preview_output is None:
if isinstance(item, tuple): if isinstance(item, tuple):
text_value = item[0] if item else '' text_value = item[0] if item else ''

View File

@ -298,6 +298,7 @@ class PreviewAudio(IO.ComfyNode):
search_aliases=["play audio"], search_aliases=["play audio"],
display_name="Preview Audio", display_name="Preview Audio",
category="audio", category="audio",
description="Preview the audio without saving it to the ComfyUI output directory.",
inputs=[ inputs=[
IO.Audio.Input("audio"), IO.Audio.Input("audio"),
], ],

View File

@ -1,3 +1,5 @@
import json
import numpy as np import numpy as np
import torch import torch
from PIL import Image, ImageDraw, ImageEnhance, ImageFont from PIL import Image, ImageDraw, ImageEnhance, ImageFont
@ -166,6 +168,111 @@ def boxes_to_regions(boxes, width: int, height: int) -> list:
return regions 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 _norm_bbox(region: dict) -> list[int]:
def grid(value: float) -> int: def grid(value: float) -> int:
return max(0, min(1000, round(value * 1000))) return max(0, min(1000, round(value * 1000)))
@ -217,29 +324,48 @@ class CreateBoundingBoxes(io.ComfyNode):
optional=True, optional=True,
tooltip="Optional image used as background in the canvas and preview.", 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, 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."), 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, 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."), tooltip="Height of the canvas and the pixel grid for the bounding boxes."),
editor_state, 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=[ outputs=[
io.Image.Output(display_name="preview"), io.Image.Output(display_name="preview"),
io.BoundingBox.Output(display_name="bboxes"), io.BoundingBox.Output(display_name="bboxes"),
io.Array.Output(display_name="elements"), io.Array.Output(display_name="elements"),
], ],
is_output_node=True,
is_experimental=True, is_experimental=True,
) )
@classmethod @classmethod
def execute(cls, width, height, editor_state=None, background=None) -> io.NodeOutput: def execute(cls, width, height, editor_state=None, last_incoming=None, background=None, bboxes=None) -> io.NodeOutput:
regions = boxes_to_regions(editor_state, width, height) 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)
preview = render_preview(regions, width, height, _bg_from_image(background)) preview = render_preview(regions, width, height, _bg_from_image(background))
ui = {"dims": [width, height]}
if incoming:
ui["input_bboxes"] = incoming
return io.NodeOutput( return io.NodeOutput(
preview, preview,
fractions_to_bbox_frame(regions, width, height), fractions_to_bbox_frame(regions, width, height),
build_elements(regions), build_elements(regions),
ui={"dims": [width, height]}, ui=ui,
) )

View File

@ -844,15 +844,18 @@ class ImageMergeTileList(IO.ComfyNode):
# Format specifications # Format specifications
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
# Maps (file_format, bit_depth, has_alpha) -> (numpy dtype scale, av pixel format, # Maps (file_format, bit_depth, num_channels) -> (quantization scale, numpy dtype,
# stream pix_fmt). Keeps the encode path declarative instead of branchy. # av frame pix_fmt, stream pix_fmt). Keeps the encode path declarative instead of branchy.
_FORMAT_SPECS = { _FORMAT_SPECS = {
("png", "8-bit", False): {"scale": 255.0, "dtype": np.uint8, "frame_fmt": "rgb24", "stream_fmt": "rgb24"}, ("png", "8-bit", 1): {"scale": 255.0, "dtype": np.uint8, "frame_fmt": "gray", "stream_fmt": "gray"},
("png", "8-bit", True): {"scale": 255.0, "dtype": np.uint8, "frame_fmt": "rgba", "stream_fmt": "rgba"}, ("png", "8-bit", 3): {"scale": 255.0, "dtype": np.uint8, "frame_fmt": "rgb24", "stream_fmt": "rgb24"},
("png", "16-bit", False): {"scale": 65535.0, "dtype": np.uint16, "frame_fmt": "rgb48le", "stream_fmt": "rgb48be"}, ("png", "8-bit", 4): {"scale": 255.0, "dtype": np.uint8, "frame_fmt": "rgba", "stream_fmt": "rgba"},
("png", "16-bit", True): {"scale": 65535.0, "dtype": np.uint16, "frame_fmt": "rgba64le", "stream_fmt": "rgba64be"}, ("png", "16-bit", 1): {"scale": 65535.0, "dtype": np.uint16, "frame_fmt": "gray16le", "stream_fmt": "gray16be"},
("exr", "32-bit float", False): {"scale": 1.0, "dtype": np.float32, "frame_fmt": "gbrpf32le", "stream_fmt": "gbrpf32le"}, ("png", "16-bit", 3): {"scale": 65535.0, "dtype": np.uint16, "frame_fmt": "rgb48le", "stream_fmt": "rgb48be"},
("exr", "32-bit float", True): {"scale": 1.0, "dtype": np.float32, "frame_fmt": "gbrapf32le", "stream_fmt": "gbrapf32le"}, ("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) return torch.cat([hlg_to_linear(rgb), alpha], dim=-1)
# Piecewise: sqrt branch below 0.5, log branch above. # 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. # values above 1.0 are allowed and extrapolate naturally.
low = (t ** 2) / 3.0 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) return torch.where(t <= 0.5, low, high)
@ -1087,7 +1091,8 @@ def _encode_image(
bit_depth: str, bit_depth: str,
colorspace: str, colorspace: str,
) -> bytes: ) -> 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 For EXR the input is interpreted according to `colorspace` and converted
to scene-linear (EXR's convention) before writing: 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 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. 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 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: if spec["dtype"] == np.float32:
# EXR path: preserve full range, no clamp. # EXR path: preserve full range, no clamp.

View File

@ -61,14 +61,10 @@ class Load3D(IO.ComfyNode):
@classmethod @classmethod
def execute(cls, model_file, image, **kwargs) -> IO.NodeOutput: 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() load_image_node = nodes.LoadImage()
output_image, ignore_mask = load_image_node.load_image(image=image_path) output_image, ignore_mask = load_image_node.load_image(image=image['image'])
ignore_image, output_mask = load_image_node.load_image(image=mask_path) ignore_image, output_mask = load_image_node.load_image(image=image['mask'])
normal_image, ignore_mask2 = load_image_node.load_image(image=normal_path) normal_image, ignore_mask2 = load_image_node.load_image(image=image['normal'])
video = None video = None
@ -96,6 +92,7 @@ class Preview3D(IO.ComfyNode):
search_aliases=["view mesh", "3d viewer"], search_aliases=["view mesh", "3d viewer"],
display_name="Preview 3D & Animation", display_name="Preview 3D & Animation",
category="3d", category="3d",
description="Preview a 3D model file without saving it to the ComfyUI output directory.",
is_experimental=True, is_experimental=True,
is_output_node=True, is_output_node=True,
inputs=[ inputs=[
@ -140,6 +137,7 @@ class Preview3DAdvanced(IO.ComfyNode):
display_name="Preview 3D (Advanced)", display_name="Preview 3D (Advanced)",
search_aliases=["preview 3d", "3d viewer", "view mesh", "frame 3d", "3d camera output"], search_aliases=["preview 3d", "3d viewer", "view mesh", "frame 3d", "3d camera output"],
category="3d", category="3d",
description="Preview a 3D model file without saving it to the ComfyUI output directory.",
is_experimental=True, is_experimental=True,
is_output_node=True, is_output_node=True,
inputs=[ inputs=[
@ -176,8 +174,9 @@ class Preview3DAdvanced(IO.ComfyNode):
filename = f"preview3d_advanced_{uuid.uuid4().hex}.{model_3d.format}" filename = f"preview3d_advanced_{uuid.uuid4().hex}.{model_3d.format}"
model_3d.save_to(os.path.join(folder_paths.get_temp_directory(), filename)) 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_input = kwargs.get("camera_info", None)
camera_info = camera_info_input if camera_info_input is not None else viewport_state['camera_info'] 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_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', []) 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( return IO.NodeOutput(
@ -197,6 +196,7 @@ class PreviewGaussianSplat(IO.ComfyNode):
node_id="PreviewGaussianSplat", node_id="PreviewGaussianSplat",
display_name="Preview Splat", display_name="Preview Splat",
category="3d", category="3d",
description="Preview a gaussian splat 3D file without saving it to the ComfyUI output directory.",
is_experimental=True, is_experimental=True,
is_output_node=True, is_output_node=True,
search_aliases=[ search_aliases=[
@ -244,8 +244,9 @@ class PreviewGaussianSplat(IO.ComfyNode):
filename = f"preview_splat_{uuid.uuid4().hex}.{model_3d.format}" filename = f"preview_splat_{uuid.uuid4().hex}.{model_3d.format}"
model_3d.save_to(os.path.join(folder_paths.get_temp_directory(), filename)) 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_input = kwargs.get("camera_info", None)
camera_info = camera_info_input if camera_info_input is not None else viewport_state['camera_info'] 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_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', []) 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( return IO.NodeOutput(
@ -265,6 +266,7 @@ class PreviewPointCloud(IO.ComfyNode):
node_id="PreviewPointCloud", node_id="PreviewPointCloud",
display_name="Preview Point Cloud", display_name="Preview Point Cloud",
category="3d", category="3d",
description="Preview a point cloud 3D file without saving it to the ComfyUI output directory.",
is_experimental=True, is_experimental=True,
is_output_node=True, is_output_node=True,
search_aliases=[ search_aliases=[
@ -303,8 +305,9 @@ class PreviewPointCloud(IO.ComfyNode):
filename = f"preview_pointcloud_{uuid.uuid4().hex}.{model_3d.format}" filename = f"preview_pointcloud_{uuid.uuid4().hex}.{model_3d.format}"
model_3d.save_to(os.path.join(folder_paths.get_temp_directory(), filename)) 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_input = kwargs.get("camera_info", None)
camera_info = camera_info_input if camera_info_input is not None else viewport_state['camera_info'] 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_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', []) 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( return IO.NodeOutput(
@ -375,8 +378,9 @@ class Load3DAdvanced(IO.ComfyNode):
file_3d = None file_3d = None
if model_file and model_file != "none": if model_file and model_file != "none":
file_3d = Types.File3D(folder_paths.get_annotated_filepath(model_file)) 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', []) model_3d_info = viewport_state.get('model_3d_info', [])
return IO.NodeOutput(file_3d, model_3d_info, viewport_state['camera_info'], width, height) return IO.NodeOutput(file_3d, model_3d_info, viewport_state.get('camera_info'), width, height)
class Load3DExtension(ComfyExtension): class Load3DExtension(ComfyExtension):

View File

@ -419,17 +419,18 @@ class MaskPreview(IO.ComfyNode):
search_aliases=["show mask", "view mask", "inspect mask", "debug mask"], search_aliases=["show mask", "view mask", "inspect mask", "debug mask"],
display_name="Preview Mask", display_name="Preview Mask",
category="image/mask", category="image/mask",
description="Saves the input images to your ComfyUI output directory.", description="Preview the masks without saving them to the ComfyUI output directory.",
inputs=[ inputs=[
IO.Mask.Input("mask"), IO.Mask.Input("mask"),
], ],
hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo], hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo],
is_output_node=True, is_output_node=True,
outputs=[IO.Mask.Output(display_name="mask")]
) )
@classmethod @classmethod
def execute(cls, mask, filename_prefix="ComfyUI") -> IO.NodeOutput: def execute(cls, mask, filename_prefix="ComfyUI") -> IO.NodeOutput:
return IO.NodeOutput(ui=UI.PreviewMask(mask)) return IO.NodeOutput(mask, ui=UI.PreviewMask(mask))
class MaskExtension(ComfyExtension): class MaskExtension(ComfyExtension):

View File

@ -18,6 +18,7 @@ class PreviewAny():
CATEGORY = "utilities" CATEGORY = "utilities"
SEARCH_ALIASES = ["show output", "inspect", "debug", "print value", "show text"] SEARCH_ALIASES = ["show output", "inspect", "debug", "print value", "show text"]
DESCRIPTION = "Preview any input value as text."
def main(self, source=None): def main(self, source=None):
torch.set_printoptions(edgeitems=6) torch.set_printoptions(edgeitems=6)

View File

@ -10,11 +10,10 @@ class String(io.ComfyNode):
return io.Schema( return io.Schema(
node_id="PrimitiveString", node_id="PrimitiveString",
search_aliases=["text", "string", "text box", "prompt"], search_aliases=["text", "string", "text box", "prompt"],
display_name="Text String (DEPRECATED)", display_name="Text",
category="utilities/primitive", category="utilities/primitive",
inputs=[io.String.Input("value")], inputs=[io.String.Input("value")],
outputs=[io.String.Output()], outputs=[io.String.Output()]
is_deprecated=True
) )
@classmethod @classmethod
@ -28,7 +27,7 @@ class StringMultiline(io.ComfyNode):
return io.Schema( return io.Schema(
node_id="PrimitiveStringMultiline", node_id="PrimitiveStringMultiline",
search_aliases=["text", "string", "text multiline", "string multiline", "text box", "prompt"], search_aliases=["text", "string", "text multiline", "string multiline", "text box", "prompt"],
display_name="Input Text", display_name="Text (Multiline)",
category="utilities/primitive", category="utilities/primitive",
essentials_category="Basics", essentials_category="Basics",
inputs=[io.String.Input("value", multiline=True)], inputs=[io.String.Input("value", multiline=True)],

View File

@ -13,7 +13,7 @@ from typing_extensions import override
import folder_paths import folder_paths
from comfy.cli_args import args from comfy.cli_args import args
from comfy_api.latest import ComfyExtension, IO, Types from comfy_api.latest import ComfyExtension, IO, Types, UI
def pack_variable_mesh_batch(vertices, faces, colors=None, uvs=None, texture=None, unlit=False): def pack_variable_mesh_batch(vertices, faces, colors=None, uvs=None, texture=None, unlit=False):
@ -406,10 +406,165 @@ class SaveGLB(IO.ComfyNode):
return IO.NodeOutput(ui={"3d": results}) 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): class Save3DExtension(ComfyExtension):
@override @override
async def get_node_list(self) -> list[type[IO.ComfyNode]]: async def get_node_list(self) -> list[type[IO.ComfyNode]]:
return [SaveGLB] return [SaveGLB, Save3DAdvanced, SaveGaussianSplat, SavePointCloud]
async def comfy_entrypoint() -> Save3DExtension: async def comfy_entrypoint() -> Save3DExtension:

View File

@ -0,0 +1,614 @@
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()

View File

@ -0,0 +1,71 @@
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()

View File

@ -81,7 +81,7 @@ class SaveVideo(io.ComfyNode):
display_name="Save Video", display_name="Save Video",
category="video", category="video",
essentials_category="Basics", essentials_category="Basics",
description="Saves the input images to your ComfyUI output directory.", description="Saves the input videos to your ComfyUI output directory.",
inputs=[ inputs=[
io.Video.Input("video", tooltip="The video to save."), 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."), 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."),

View File

@ -1,3 +1,3 @@
# This file is automatically generated by the build process when version is # This file is automatically generated by the build process when version is
# updated in pyproject.toml. # updated in pyproject.toml.
__version__ = "0.27.0" __version__ = "0.28.0"

View File

@ -29,6 +29,7 @@ from comfy_execution.caching import (
HierarchicalCache, HierarchicalCache,
LRUCache, LRUCache,
RAMPressureCache, RAMPressureCache,
RAM_CACHE_LARGE_INTERMEDIATE,
) )
from comfy_execution.graph import ( from comfy_execution.graph import (
DynamicPrompt, DynamicPrompt,
@ -425,12 +426,12 @@ def _is_intermediate_output(dynprompt, node_id):
def _send_cached_ui(server, node_id, display_node_id, cached, prompt_id, ui_outputs): def _send_cached_ui(server, node_id, display_node_id, cached, prompt_id, ui_outputs):
if cached.ui is not None:
ui_outputs[node_id] = cached.ui
if server.client_id is None: if server.client_id is None:
return return
cached_ui = cached.ui or {} cached_ui = cached.ui or {}
server.send_sync("executed", { "node": node_id, "display_node": display_node_id, "output": cached_ui.get("output", None), "prompt_id": prompt_id }, server.client_id) server.send_sync("executed", { "node": node_id, "display_node": display_node_id, "output": cached_ui.get("output", None), "prompt_id": prompt_id }, server.client_id)
if cached.ui is not None:
ui_outputs[node_id] = cached.ui
async def execute(server, dynprompt, caches, current_item, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes, ui_outputs): async def execute(server, dynprompt, caches, current_item, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes, ui_outputs):
unique_id = current_item unique_id = current_item
@ -794,12 +795,16 @@ class PromptExecutor:
if self.cache_type == CacheType.RAM_PRESSURE: if self.cache_type == CacheType.RAM_PRESSURE:
ram_release_callback(ram_inactive_headroom) ram_release_callback(ram_inactive_headroom)
ram_shortfall = ram_headroom - psutil.virtual_memory().available ram_shortfall = ram_headroom - psutil.virtual_memory().available
freed = comfy.model_management.free_pins(ram_shortfall + 512 * (1024 ** 2)) if ram_shortfall > 0:
if freed < ram_shortfall: freed = ram_release_callback(ram_headroom, free_active=True, min_entry_size=RAM_CACHE_LARGE_INTERMEDIATE)
if freed > 64 * (1024 ** 2): ram_shortfall -= freed
# AIMDO MEM_DECOMMIT can outrun psutil.available catching up. if comfy.model_management.should_free_pins_for_ram_pressure(ram_shortfall):
time.sleep(0.05) freed = comfy.model_management.free_pins(ram_shortfall + 512 * (1024 ** 2))
ram_release_callback(ram_headroom, free_active=True) 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: else:
# Only execute when the while-loop ends without break # Only execute when the while-loop ends without break
# Send cached UI for intermediate output nodes that weren't executed # Send cached UI for intermediate output nodes that weren't executed

View File

@ -1709,6 +1709,7 @@ class PreviewImage(SaveImage):
self.compress_level = 1 self.compress_level = 1
SEARCH_ALIASES = ["preview", "preview image", "show image", "view image", "display image", "image viewer"] 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 @classmethod
def INPUT_TYPES(s): def INPUT_TYPES(s):
@ -2458,6 +2459,7 @@ async def init_builtin_extra_nodes():
"nodes_camera_trajectory.py", "nodes_camera_trajectory.py",
"nodes_edit_model.py", "nodes_edit_model.py",
"nodes_tcfg.py", "nodes_tcfg.py",
"nodes_seedvr.py",
"nodes_context_windows.py", "nodes_context_windows.py",
"nodes_qwen.py", "nodes_qwen.py",
"nodes_boogu.py", "nodes_boogu.py",
@ -2503,6 +2505,7 @@ async def init_builtin_extra_nodes():
"nodes_triposplat.py", "nodes_triposplat.py",
"nodes_depth_anything_3.py", "nodes_depth_anything_3.py",
"nodes_seed.py", "nodes_seed.py",
"nodes_text.py",
] ]
import_failed = [] import_failed = []

View File

@ -7,18 +7,18 @@ components:
description: Timestamp when the asset was created description: Timestamp when the asset was created
format: date-time format: date-time
type: string 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: hash:
description: Blake3 hash of the asset content. description: Blake3 hash of the asset content.
pattern: ^blake3:[a-f0-9]{64}$ pattern: ^blake3:[a-f0-9]{64}$
type: string 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: id:
description: Unique identifier for the asset description: Unique identifier for the asset
format: uuid format: uuid
@ -144,6 +144,14 @@ components:
AssetUpdated: AssetUpdated:
description: Response returned when an existing asset is successfully updated. description: Response returned when an existing asset is successfully updated.
properties: 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: hash:
description: Blake3 hash of the asset content. description: Blake3 hash of the asset content.
pattern: ^blake3:[a-f0-9]{64}$ pattern: ^blake3:[a-f0-9]{64}$
@ -1636,7 +1644,7 @@ paths:
format: uuid format: uuid
type: string type: string
tags: 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 type: string
user_metadata: user_metadata:
description: Custom JSON metadata as a string description: Custom JSON metadata as a string
@ -1821,7 +1829,7 @@ paths:
content: content:
application/json: application/json:
schema: schema:
$ref: '#/components/schemas/Asset' $ref: '#/components/schemas/AssetUpdated'
description: Asset updated successfully description: Asset updated successfully
"400": "400":
content: content:
@ -2462,9 +2470,6 @@ paths:
supports_preview_metadata: supports_preview_metadata:
description: Whether the server supports preview metadata description: Whether the server supports preview metadata
type: boolean type: boolean
supports_model_type_tags:
description: Whether the server supports namespaced model type asset tags
type: boolean
type: object type: object
description: Success description: Success
headers: headers:
@ -3292,6 +3297,12 @@ paths:
schema: schema:
$ref: '#/components/schemas/ErrorResponse' $ref: '#/components/schemas/ErrorResponse'
description: Invalid request parameters description: Invalid request parameters
"401":
content:
application/json:
schema:
$ref: '#/components/schemas/ErrorResponse'
description: Unauthorized - Authentication required
"500": "500":
content: content:
application/json: application/json:

View File

@ -1,6 +1,6 @@
[project] [project]
name = "ComfyUI" name = "ComfyUI"
version = "0.27.0" version = "0.28.0"
readme = "README.md" readme = "README.md"
license = { file = "LICENSE" } license = { file = "LICENSE" }
requires-python = ">=3.10" requires-python = ">=3.10"

View File

@ -1,6 +1,6 @@
comfyui-frontend-package==1.45.20 comfyui-frontend-package==1.45.21
comfyui-workflow-templates==0.11.6 comfyui-workflow-templates==0.11.9
comfyui-embedded-docs==0.5.7 comfyui-embedded-docs==0.5.8
torch torch
torchsde torchsde
torchvision torchvision
@ -22,7 +22,7 @@ alembic
SQLAlchemy>=2.0.0 SQLAlchemy>=2.0.0
filelock filelock
av>=16.0.0 av>=16.0.0
comfy-kitchen==0.2.16 comfy-kitchen==0.2.20
comfy-aimdo==0.4.10 comfy-aimdo==0.4.10
requests requests
simpleeval>=1.0.0 simpleeval>=1.0.0

View File

@ -24,6 +24,28 @@ def app(model_manager):
app.add_routes(routes) app.add_routes(routes)
return app 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): async def test_get_model_preview_safetensors(aiohttp_client, app, tmp_path):
img = Image.new('RGB', (100, 100), 'white') img = Image.new('RGB', (100, 100), 'white')
img_byte_arr = BytesIO() img_byte_arr = BytesIO()

View File

@ -0,0 +1,186 @@
"""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()

View File

@ -0,0 +1,55 @@
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)

View File

@ -0,0 +1,51 @@
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)

View File

@ -0,0 +1,77 @@
"""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)

View File

@ -15,7 +15,7 @@ if not has_gpu():
args.cpu = True args.cpu = True
from comfy import ops from comfy import ops
from comfy.quant_ops import QuantizedTensor from comfy.quant_ops import QUANT_ALGOS, QuantizedTensor
import comfy.utils import comfy.utils
@ -283,7 +283,59 @@ class TestMixedPrecisionOps(unittest.TestCase):
saved = model.state_dict() saved = model.state_dict()
saved_conf = json.loads(saved["layer.comfy_quant"].numpy().tobytes()) saved_conf = json.loads(saved["layer.comfy_quant"].numpy().tobytes())
self.assertTrue(saved_conf["convrot"]) 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["convrot_groupsize"], 256)
self.assertEqual(saved_conf["linear_dtype"], "int8")
self.assertNotIn("quant_group_size", saved_conf)
if __name__ == "__main__": if __name__ == "__main__":
unittest.main() unittest.main()

View File

@ -2,7 +2,7 @@ from collections import defaultdict
import torch import torch
from comfy.model_detection import detect_unet_config, model_config_from_unet_config from comfy.model_detection import detect_unet_config, model_config_from_unet, model_config_from_unet_config
import comfy.supported_models import comfy.supported_models
@ -73,6 +73,49 @@ def _make_flux_schnell_comfyui_sd():
return 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: class TestModelDetection:
"""Verify that first-match model detection selects the correct model """Verify that first-match model detection selects the correct model
based on list ordering and unet_config specificity.""" based on list ordering and unet_config specificity."""
@ -125,6 +168,96 @@ class TestModelDetection:
assert model_config is not None assert model_config is not None
assert type(model_config).__name__ == "FluxSchnell" 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): def test_unet_config_and_required_keys_combination_is_unique(self):
"""Each model in the registry must have a unique combination of """Each model in the registry must have a unique combination of
``unet_config`` and ``required_keys``. If two models share the same ``unet_config`` and ``required_keys``. If two models share the same

View File

@ -0,0 +1,74 @@
"""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")

View File

@ -0,0 +1,79 @@
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

View File

@ -0,0 +1,169 @@
"""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)

View File

@ -0,0 +1,320 @@
"""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,
},
}
]

View File

@ -0,0 +1,94 @@
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)

View File

@ -0,0 +1,407 @@
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)

View File

@ -818,6 +818,30 @@ class TestExecution:
except urllib.error.HTTPError: except urllib.error.HTTPError:
pass # Expected behavior pass # Expected behavior
def test_cached_outputs_in_job_without_client_id(self, client: ComfyClient, builder: GraphBuilder):
g = builder
image = g.node("StubImage", content="BLACK", height=32, width=32, batch_size=1)
output = g.node("SaveImage", images=image.out(0))
# Prime the cache with a normal run.
client.run(g)
# Resubmit anonymously (no client_id) so output nodes are cache hits with no websocket client.
data = json.dumps({"prompt": g.finalize()}).encode('utf-8')
req = urllib.request.Request(f"http://{client.server_address}/prompt", data=data)
prompt_id = json.loads(urllib.request.urlopen(req).read())['prompt_id']
for _ in range(100):
job = client.get_job(prompt_id)
if job is not None and job['status'] not in ('pending', 'in_progress'):
break
time.sleep(0.1)
else:
raise AssertionError("Prompt did not complete in time")
assert job['status'] == 'completed'
assert output.id in job['outputs'], "Cached outputs must appear in job outputs without a client_id"
def _create_history_item(self, client, builder): def _create_history_item(self, client, builder):
g = GraphBuilder(prefix="offset_test") g = GraphBuilder(prefix="offset_test")
input_node = g.node( input_node = g.node(