[pull] master from Comfy-Org:master by pull[bot] · Pull Request #1152 · code/app-python-comfyui · GitHub
Skip to content

[pull] master from Comfy-Org:master #1152

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
pull merged 1 commit into from
Sep 28, 2026
Merged

[pull] master from Comfy-Org:master #1152

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 8 additions & 4 deletions comfy/ldm/audio/dit.py
View file
Open in desktop
6 changes: 4 additions & 2 deletions comfy/ldm/cascade/common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,12 +18,13 @@

import torch
import torch.nn as nn
from comfy.ldm.modules.attention import optimized_attention
from comfy.ldm.modules.attention import AttentionTensorContainer, ComfyAttention, optimized_attention
import comfy.ops

class OptimizedAttention(nn.Module):
def __init__(self, c, nhead, dropout=0.0, dtype=None, device=None, operations=None):
super().__init__()
self.comfy_attention = ComfyAttention()
self.heads = nhead

self.to_q = operations.Linear(c, c, bias=True, dtype=dtype, device=device)
Expand All @@ -37,7 +38,8 @@ def forward(self, q, k, v, transformer_options={}):
k = self.to_k(k)
v = self.to_v(v)

out = optimized_attention(q, k, v, self.heads, transformer_options=transformer_options)
q, k, v = AttentionTensorContainer(q), AttentionTensorContainer(k), AttentionTensorContainer(v)
out = optimized_attention(q, k, v, self.heads, preferred_attention=self.comfy_attention, transformer_options=transformer_options)

return self.out_proj(out)

Expand Down
10 changes: 7 additions & 3 deletions comfy/ldm/cogvideo/model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
import torch.nn as nn
import torch.nn.functional as F

from comfy.ldm.modules.attention import optimized_attention
from comfy.ldm.modules.attention import AttentionTensorContainer, ComfyAttention, optimized_attention
import comfy.patcher_extension
import comfy.ldm.common_dit

Expand Down Expand Up @@ -241,6 +241,7 @@ def __init__(self, dim, num_heads, head_dim, time_dim,
eps=1e-5, ff_inner_dim=None, ff_bias=True,
device=None, dtype=None, operations=None):
super().__init__()
self.comfy_attention = ComfyAttention()
self.dim = dim
self.num_heads = num_heads
self.head_dim = head_dim
Expand Down Expand Up @@ -290,12 +291,15 @@ def forward(self, hidden_states, encoder_hidden_states, temb, image_rotary_emb=N
k_img = apply_rotary_emb(k_img, image_rotary_emb)
q = torch.cat([q[:, :text_seq_length], q_img.transpose(1, 2)], dim=1)
k = torch.cat([k[:, :text_seq_length], k_img.transpose(1, 2)], dim=1)
del q_img, k_img

q, k, v = AttentionTensorContainer(q.reshape(b, s, n * d)), AttentionTensorContainer(k.reshape(b, s, n * d)), AttentionTensorContainer(v)
attn_out = optimized_attention(
q.reshape(b, s, n * d),
k.reshape(b, s, n * d),
q,
k,
v,
heads=self.num_heads,
preferred_attention=self.comfy_attention,
transformer_options=transformer_options,
)

Expand Down
33 changes: 20 additions & 13 deletions comfy/ldm/cosmos/blocks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,9 @@
from einops.layers.torch import Rearrange
from torch import nn

from comfy.ldm.modules.attention import optimized_attention
from comfy.ldm.modules.attention import AttentionTensorContainer, ComfyAttention, optimized_attention
import comfy.model_management
import comfy.quant_ops


def get_normalization(name: str, channels: int, weight_args={}, operations=None):
Expand Down Expand Up @@ -89,6 +91,7 @@ def __init__(
operations=None,
) -> None:
super().__init__()
self.comfy_attention = ComfyAttention()

self.is_selfattn = context_dim is None # self attention

Expand Down Expand Up @@ -157,17 +160,20 @@ def cal_qkv(
k = self.to_k[1](k)
v = self.to_v[1](v)
if self.is_selfattn and rope_emb is not None: # only apply to self-attention!
# apply_rotary_pos_emb inlined
q_shape = q.shape
q = q.reshape(*q.shape[:-1], 2, -1).movedim(-2, -1).unsqueeze(-2)
q = rope_emb[..., 0] * q[..., 0] + rope_emb[..., 1] * q[..., 1]
q = q.movedim(-1, -2).reshape(*q_shape).to(x.dtype)

# apply_rotary_pos_emb inlined
k_shape = k.shape
k = k.reshape(*k.shape[:-1], 2, -1).movedim(-2, -1).unsqueeze(-2)
k = rope_emb[..., 0] * k[..., 0] + rope_emb[..., 1] * k[..., 1]
k = k.movedim(-1, -2).reshape(*k_shape).to(x.dtype)
if not comfy.model_management.in_training:
q, k = comfy.quant_ops.ck.apply_rope_split_half(q, k, rope_emb)
else:
# apply_rotary_pos_emb inlined
q_shape = q.shape
q = q.reshape(*q.shape[:-1], 2, -1).movedim(-2, -1).unsqueeze(-2)
q = rope_emb[..., 0] * q[..., 0] + rope_emb[..., 1] * q[..., 1]
q = q.movedim(-1, -2).reshape(*q_shape).to(x.dtype)

# apply_rotary_pos_emb inlined
k_shape = k.shape
k = k.reshape(*k.shape[:-1], 2, -1).movedim(-2, -1).unsqueeze(-2)
k = rope_emb[..., 0] * k[..., 0] + rope_emb[..., 1] * k[..., 1]
k = k.movedim(-1, -2).reshape(*k_shape).to(x.dtype)
return q, k, v

def forward(
Expand All @@ -185,7 +191,8 @@ def forward(
context (Optional[Tensor]): The key tensor of shape [B, Mk, K] or use x as context [self attention] if None
"""
q, k, v = self.cal_qkv(x, context, mask, rope_emb=rope_emb, **kwargs)
out = optimized_attention(q, k, v, self.heads, skip_reshape=True, mask=mask, skip_output_reshape=True, transformer_options=transformer_options)
q, k, v = AttentionTensorContainer(q), AttentionTensorContainer(k), AttentionTensorContainer(v)
out = optimized_attention(q, k, v, self.heads, skip_reshape=True, mask=mask, skip_output_reshape=True, preferred_attention=self.comfy_attention, transformer_options=transformer_options)
del q, k, v
out = rearrange(out, " b n s c -> s b (n c)")
return self.to_out(out)
Expand Down
14 changes: 10 additions & 4 deletions comfy/ldm/cosmos/predict2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
from torchvision import transforms

import comfy.patcher_extension
from comfy.ldm.modules.attention import optimized_attention
from comfy.ldm.modules.attention import AttentionTensorContainer, ComfyAttention, optimized_attention
import comfy.ldm.common_dit
import comfy.ops
import comfy.quant_ops
Expand All @@ -38,7 +38,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
return x


def torch_attention_op(q_B_S_H_D: torch.Tensor, k_B_S_H_D: torch.Tensor, v_B_S_H_D: torch.Tensor, transformer_options: Optional[dict] = {}) -> torch.Tensor:
def torch_attention_op(q_B_S_H_D: torch.Tensor, k_B_S_H_D: torch.Tensor, v_B_S_H_D: torch.Tensor, transformer_options: Optional[dict] = {}, preferred_attention=None) -> torch.Tensor:
"""Computes multi-head attention using PyTorch's native implementation.

This function provides a PyTorch backend alternative to Transformer Engine's attention operation.
Expand All @@ -60,12 +60,16 @@ def torch_attention_op(q_B_S_H_D: torch.Tensor, k_B_S_H_D: torch.Tensor, v_B_S_H
Returns:
Attention output tensor with shape (batch, seq_len, n_heads * head_dim)
"""
if isinstance(q_B_S_H_D, AttentionTensorContainer):
q_B_S_H_D, k_B_S_H_D, v_B_S_H_D = q_B_S_H_D.take(), k_B_S_H_D.take(), v_B_S_H_D.take()
in_q_shape = q_B_S_H_D.shape
in_k_shape = k_B_S_H_D.shape
q_B_H_S_D = rearrange(q_B_S_H_D, "b ... h k -> b h ... k").view(in_q_shape[0], in_q_shape[-2], -1, in_q_shape[-1])
k_B_H_S_D = rearrange(k_B_S_H_D, "b ... h v -> b h ... v").view(in_k_shape[0], in_k_shape[-2], -1, in_k_shape[-1])
v_B_H_S_D = rearrange(v_B_S_H_D, "b ... h v -> b h ... v").view(in_k_shape[0], in_k_shape[-2], -1, in_k_shape[-1])
return optimized_attention(q_B_H_S_D, k_B_H_S_D, v_B_H_S_D, in_q_shape[-2], skip_reshape=True, transformer_options=transformer_options)
q_B_H_S_D, k_B_H_S_D, v_B_H_S_D = AttentionTensorContainer(q_B_H_S_D), AttentionTensorContainer(k_B_H_S_D), AttentionTensorContainer(v_B_H_S_D)
del q_B_S_H_D, k_B_S_H_D, v_B_S_H_D
return optimized_attention(q_B_H_S_D, k_B_H_S_D, v_B_H_S_D, in_q_shape[-2], skip_reshape=True, preferred_attention=preferred_attention, transformer_options=transformer_options)


class Attention(nn.Module):
Expand Down Expand Up @@ -112,6 +116,7 @@ def __init__(
operations=None,
) -> None:
super().__init__()
self.comfy_attention = ComfyAttention()
logging.debug(
f"Setting up {self.__class__.__name__}. Query dim is {query_dim}, context_dim is {context_dim} and using "
f"{n_heads} heads with a dimension of {head_dim}."
Expand Down Expand Up @@ -197,7 +202,7 @@ def apply_norm_and_rotary_pos_emb(
return q, k, v

def compute_attention(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, transformer_options: Optional[dict] = {}) -> torch.Tensor:
result = self.attn_op(q, k, v, transformer_options=transformer_options) # [B, S, H, D]
result = self.attn_op(q, k, v, transformer_options=transformer_options, preferred_attention=self.comfy_attention) # [B, S, H, D]
return self.output_dropout(self.output_proj(result))

def forward(
Expand All @@ -213,6 +218,7 @@ def forward(
context (Optional[Tensor]): The key tensor of shape [B, Mk, K] or use x as context [self attention] if None
"""
q, k, v = self.compute_qkv(x, context, rope_emb=rope_emb, transformer_options=transformer_options)
q, k, v = AttentionTensorContainer(q), AttentionTensorContainer(k), AttentionTensorContainer(v)
return self.compute_attention(q, k, v, transformer_options=transformer_options)


Expand Down
12 changes: 8 additions & 4 deletions comfy/ldm/genmo/joint_model/asymm_models_joint.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
import torch.nn.functional as F
from einops import rearrange
# from flash_attn import flash_attn_varlen_qkvpacked_func
from comfy.ldm.modules.attention import optimized_attention
from comfy.ldm.modules.attention import AttentionTensorContainer, ComfyAttention, optimized_attention

from .layers import (
FeedForward,
Expand Down Expand Up @@ -68,6 +68,7 @@ def __init__(
operations=None,
):
super().__init__()
self.comfy_attention = ComfyAttention()
self.dim_x = dim_x
self.dim_y = dim_y
self.num_heads = num_heads
Expand Down Expand Up @@ -142,13 +143,16 @@ def forward(
k = torch.cat([k_x, k_y[:, :crop_y]], dim=1).transpose(1, 2)
v = torch.cat([v_x, v_y[:, :crop_y]], dim=1).transpose(1, 2)

seq_x, seq_y = q_x.shape[1], q_y.shape[1]
del q_x, k_x, v_x, q_y, k_y, v_y
q, k, v = AttentionTensorContainer(q), AttentionTensorContainer(k), AttentionTensorContainer(v)
xy = optimized_attention(q,
k,
v, self.num_heads, skip_reshape=True, transformer_options=transformer_options)
v, self.num_heads, skip_reshape=True, preferred_attention=self.comfy_attention, transformer_options=transformer_options)

x, y = torch.tensor_split(xy, (q_x.shape[1],), dim=1)
x, y = torch.tensor_split(xy, (seq_x,), dim=1)
x = self.proj_x(x)
o = torch.zeros(y.shape[0], q_y.shape[1], y.shape[-1], device=y.device, dtype=y.dtype)
o = torch.zeros(y.shape[0], seq_y, y.shape[-1], device=y.device, dtype=y.dtype)
o[:, :y.shape[1]] = y

y = self.proj_y(o)
Expand Down
8 changes: 5 additions & 3 deletions comfy/ldm/genmo/joint_model/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
from typing import Optional

import torch
from comfy.ldm.modules.attention import AttentionTensorContainer, ComfyAttention, optimized_attention
import torch.nn as nn
import torch.nn.functional as F

Expand Down Expand Up @@ -51,6 +52,7 @@ def __init__(
output_dim (int): Dimensionality of output tokens. Defaults to embed_dim.
"""
super().__init__()
self.comfy_attention = ComfyAttention()
self.num_heads = num_heads
self.to_kv = operations.Linear(embed_dim, 2 * embed_dim, device=device, dtype=dtype)
self.to_q = operations.Linear(embed_dim, embed_dim, device=device, dtype=dtype)
Expand Down Expand Up @@ -92,9 +94,9 @@ def forward(self, x, mask):
q = q.unsqueeze(2) # (B, H, 1, head_dim)

# Compute attention.
x = F.scaled_dot_product_attention(
q, k, v, attn_mask=attn_mask, dropout_p=0.0
) # (B, H, 1, head_dim)
del kv
q, k, v = AttentionTensorContainer(q), AttentionTensorContainer(k), AttentionTensorContainer(v)
x = optimized_attention(q, k, v, self.num_heads, mask=attn_mask, skip_reshape=True, skip_output_reshape=True, preferred_attention=self.comfy_attention)

# Concatenate heads and run output.
x = x.squeeze(2).flatten(1, 2) # (B, D = H * head_dim)
Expand Down
7 changes: 5 additions & 2 deletions comfy/ldm/hunyuan_video/model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
import comfy.patcher_extension
import comfy.ldm.flux.layers
import comfy.ldm.modules.diffusionmodules.mmdit
from comfy.ldm.modules.attention import optimized_attention
from comfy.ldm.modules.attention import AttentionTensorContainer, ComfyAttention, optimized_attention

from dataclasses import dataclass
from einops import repeat
Expand Down Expand Up @@ -63,6 +63,7 @@ def __init__(
operations=None
):
super().__init__()
self.comfy_attention = ComfyAttention()
self.heads = heads
mlp_hidden_dim = hidden_size * 4

Expand All @@ -88,7 +89,9 @@ def forward(self, x, c, mask, transformer_options={}):
norm_x = self.norm1(x)
qkv = self.self_attn.qkv(norm_x)
q, k, v = qkv.reshape(qkv.shape[0], qkv.shape[1], 3, self.heads, -1).permute(2, 0, 3, 1, 4)
attn = optimized_attention(q, k, v, self.heads, mask=mask, skip_reshape=True, transformer_options=transformer_options)
q, k, v = AttentionTensorContainer(q), AttentionTensorContainer(k), AttentionTensorContainer(v)
del qkv
attn = optimized_attention(q, k, v, self.heads, mask=mask, skip_reshape=True, preferred_attention=self.comfy_attention, transformer_options=transformer_options)

x = x + self.self_attn.proj(attn) * mod1.unsqueeze(1)
x = x + self.mlp(self.norm2(x)) * mod2.unsqueeze(1)
Expand Down
12 changes: 7 additions & 5 deletions comfy/ldm/ideogram4/model.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
import comfy.patcher_extension
import comfy.quant_ops
from comfy.ldm.lumina.model import FeedForward
from comfy.ldm.modules.attention import optimized_attention_masked
from comfy.ldm.modules.attention import AttentionTensorContainer, ComfyAttention, optimized_attention_masked
from comfy.text_encoders.llama import precompute_freqs_cis

# Per-token role indicators
Expand Down Expand Up @@ -47,6 +47,7 @@ def _apply_rope_split_half1(x, freqs_cis):
class Ideogram4Attention(nn.Module):
def __init__(self, hidden_size, num_heads, eps=1e-5, dtype=None, device=None, operations=None):
super().__init__()
self.comfy_attention = ComfyAttention()
self.num_heads = num_heads
self.head_dim = hidden_size // num_heads
self.hidden_size = hidden_size
Expand Down Expand Up @@ -74,11 +75,12 @@ def forward(self, x, attn_mask, freqs_cis, transformer_options={}):
comfy.ops.uncast_bias_weight(self.norm_k, k_scale, None, k_offload_stream)

# (B, heads, L, head_dim)
q = q.transpose(1, 2)
k = k.transpose(1, 2)
v = v.transpose(1, 2)
q = AttentionTensorContainer(q.transpose(1, 2))
k = AttentionTensorContainer(k.transpose(1, 2))
v = AttentionTensorContainer(v.transpose(1, 2))
del qkv

out = optimized_attention_masked(q, k, v, self.num_heads, attn_mask, skip_reshape=True, transformer_options=transformer_options)
out = optimized_attention_masked(q, k, v, self.num_heads, attn_mask, skip_reshape=True, preferred_attention=self.comfy_attention, transformer_options=transformer_options)
return self.o(out)


Expand Down
12 changes: 7 additions & 5 deletions comfy/ldm/joyimage/model.py
View file
Open in desktop
Loading
Loading