RomanNum3ral
/
ComfyUI
mirror of https://github.com/Comfy-Org/ComfyUI.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Merge branch 'master' into cloud-openapi-projection

This commit is contained in:
Alexis Rolland 2026-07-12 13:00:15 +08:00 committed by GitHub
parent ccf32f315f 8b099de36a
commit 66dd23ec81
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
12 changed files with 79 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
comfy/ldm/modules/attention.py
View File

@ -709,7 +709,7 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
return out
try:
@torch.library.custom_op("flash_attention::flash_attn", mutates_args=())
@torch.library.custom_op("comfy::flash_attn", mutates_args=())
def flash_attn_wrapper(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
dropout_p: float = 0.0, causal: bool = False, softmax_scale: float = -1.0) -> torch.Tensor:
softmax_scale_arg = None if softmax_scale == -1.0 else softmax_scale

20
comfy/ldm/seedvr/vae.py
View File

@ -30,7 +30,7 @@ from enum import Enum
import logging
import comfy.model_management
import comfy.ops
ops = comfy.ops.disable_weight_init
ops = comfy.ops.manual_cast
def _seedvr2_temporal_slicing_min_size(temporal_size, temporal_overlap, temporal_scale=1):
@ -103,11 +103,10 @@ def tiled_vae(
storage_device = vae_model.device
result = None
count = None
def run_temporal_chunks(spatial_tile, model=vae_model, device=storage_device):
device = torch.device(device)
t_chunk = spatial_tile.to(device=device, dtype=next(model.parameters()).dtype, non_blocking=True).contiguous()
def run_temporal_chunks(spatial_tile, model=vae_model):
t_chunk = spatial_tile.contiguous()
old_device = getattr(model, "device", None)
model.device = device
model.device = t_chunk.device
old_slicing_min_size = getattr(model, slicing_attr, None)
if old_slicing_min_size is not None and slicing_min_size is not None:
if slicing_min_size <= 0:
@ -397,7 +396,7 @@ class Attention(nn.Module):
def causal_norm_wrapper(norm_layer: nn.Module, x: torch.Tensor) -> torch.Tensor:
input_dtype = x.dtype
if isinstance(norm_layer, (ops.LayerNorm, ops.RMSNorm)):
if isinstance(norm_layer, (nn.LayerNorm, nn.RMSNorm)):
if x.ndim == 4:
x = x.permute(0, 2, 3, 1)
x = norm_layer(x)
@ -408,14 +407,14 @@ def causal_norm_wrapper(norm_layer: nn.Module, x: torch.Tensor) -> torch.Tensor:
x = norm_layer(x)
x = x.permute(0, 4, 1, 2, 3)
return x.to(input_dtype)
if isinstance(norm_layer, (ops.GroupNorm, nn.BatchNorm2d, nn.SyncBatchNorm)):
if isinstance(norm_layer, (nn.GroupNorm, nn.BatchNorm2d, nn.SyncBatchNorm)):
if x.ndim <= 4:
return norm_layer(x).to(input_dtype)
if x.ndim == 5:
b, c, t, h, w = x.shape
x = x.transpose(1, 2).reshape(b * t, c, h, w)
memory_occupy = x.numel() * x.element_size() / 1024**3
if isinstance(norm_layer, ops.GroupNorm) and memory_occupy > get_norm_limit():
if isinstance(norm_layer, nn.GroupNorm) and memory_occupy > get_norm_limit():
num_chunks = min(BYTEDANCE_GN_CHUNKS_FP16 if x.element_size() == 2 else BYTEDANCE_GN_CHUNKS_FP32, norm_layer.num_groups)
if norm_layer.num_groups % num_chunks != 0:
raise ValueError(
@ -423,9 +422,9 @@ def causal_norm_wrapper(norm_layer: nn.Module, x: torch.Tensor) -> torch.Tensor:
)
num_groups_per_chunk = norm_layer.num_groups // num_chunks
weights = comfy.ops.cast_to_input(norm_layer.weight, x).chunk(num_chunks, dim=0)
biases = comfy.ops.cast_to_input(norm_layer.bias, x).chunk(num_chunks, dim=0)
x = list(x.chunk(num_chunks, dim=1))
weights = norm_layer.weight.chunk(num_chunks, dim=0)
biases = norm_layer.bias.chunk(num_chunks, dim=0)
for i, (w, bias) in enumerate(zip(weights, biases)):
x[i] = F.group_norm(x[i], num_groups_per_chunk, w, bias, norm_layer.eps)
x[i] = x[i].to(input_dtype)
@ -1459,7 +1458,6 @@ class VideoAutoencoderKLWrapper(VideoAutoencoderKL):
def _encode_with_raw_latent(self, x):
if x.ndim == 4:
x = x.unsqueeze(2)
x = x.to(dtype=next(self.parameters()).dtype)
self.device = x.device
p = super().encode(x)
z = p.squeeze(2)

2
comfy/quant_ops.py
View File

@ -48,7 +48,7 @@ try:
# 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()):
elif args.enable_triton_backend: # or (torch.version.hip is not None and _rocm_kitchen_arch_supported()):
try:
import triton
triton_version = tuple(int(v) for v in triton.__version__.split(".")[:2])

1
comfy_extras/nodes_audio.py
View File

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

4
comfy_extras/nodes_load_3d.py
View File

@ -92,6 +92,7 @@ class Preview3D(IO.ComfyNode):
search_aliases=["view mesh", "3d viewer"],
display_name="Preview 3D & Animation",
category="3d",
description="Preview a 3D model file without saving it to the ComfyUI output directory.",
is_experimental=True,
is_output_node=True,
inputs=[
@ -136,6 +137,7 @@ class Preview3DAdvanced(IO.ComfyNode):
display_name="Preview 3D (Advanced)",
search_aliases=["preview 3d", "3d viewer", "view mesh", "frame 3d", "3d camera output"],
category="3d",
description="Preview a 3D model file without saving it to the ComfyUI output directory.",
is_experimental=True,
is_output_node=True,
inputs=[
@ -193,6 +195,7 @@ class PreviewGaussianSplat(IO.ComfyNode):
node_id="PreviewGaussianSplat",
display_name="Preview Splat",
category="3d",
description="Preview a gaussian splat 3D file without saving it to the ComfyUI output directory.",
is_experimental=True,
is_output_node=True,
search_aliases=[
@ -261,6 +264,7 @@ class PreviewPointCloud(IO.ComfyNode):
node_id="PreviewPointCloud",
display_name="Preview Point Cloud",
category="3d",
description="Preview a point cloud 3D file without saving it to the ComfyUI output directory.",
is_experimental=True,
is_output_node=True,
search_aliases=[

5
comfy_extras/nodes_mask.py
View File

@ -419,17 +419,18 @@ class MaskPreview(IO.ComfyNode):
search_aliases=["show mask", "view mask", "inspect mask", "debug mask"],
display_name="Preview 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=[
IO.Mask.Input("mask"),
],
hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo],
is_output_node=True,
outputs=[IO.Mask.Output(display_name="mask")]
)
@classmethod
def execute(cls, mask, filename_prefix="ComfyUI") -> IO.NodeOutput:
return IO.NodeOutput(ui=UI.PreviewMask(mask))
return IO.NodeOutput(mask, ui=UI.PreviewMask(mask))
class MaskExtension(ComfyExtension):

1
comfy_extras/nodes_preview_any.py
View File

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

7
comfy_extras/nodes_primitive.py
View File

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

2
comfy_extras/nodes_video.py
View File

@ -81,7 +81,7 @@ class SaveVideo(io.ComfyNode):
display_name="Save Video",
category="video",
essentials_category="Basics",
description="Saves the input images to your ComfyUI output directory.",
description="Saves the input videos to your ComfyUI output directory.",
inputs=[
io.Video.Input("video", tooltip="The video to save."),
io.String.Input("filename_prefix", default="video/ComfyUI", tooltip="The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."),

1
nodes.py
View File

@ -1709,6 +1709,7 @@ class PreviewImage(SaveImage):
self.compress_level = 1
SEARCH_ALIASES = ["preview", "preview image", "show image", "view image", "display image", "image viewer"]
DESCRIPTION = "Preview the images without saving them to the ComfyUI output directory."
@classmethod
def INPUT_TYPES(s):

29
tests-unit/comfy_test/test_seedvr2_dtype.py
View File

@ -1,4 +1,5 @@
import torch
import torch.nn as nn
from comfy.cli_args import args as cli_args
@ -48,3 +49,31 @@ def test_seedvr2_vae_decode_memory_covers_full_frame_lab_transfer():
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

25
tests-unit/comfy_test/test_seedvr2_vae_tiled.py
View File

@ -122,6 +122,31 @@ def test_tiled_vae_encode_uses_tensor_return_without_indexing():
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):