Add SeedVR2 support (CORE-6) (#14424) · code/app-python-comfyui@8e2e54e · GitHub
Skip to content

Commit 8e2e54e

Browse files
authored
Add SeedVR2 support (CORE-6) (Comfy-Org#14424)
1 parent e2a6e30 commit 8e2e54e

26 files changed

Lines changed: 5712 additions & 29 deletions

‎comfy/latent_formats.py‎

Lines changed: 4 additions & 0 deletions

‎comfy/ldm/modules/diffusionmodules/model.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ def torch_cat_if_needed(xl, dim):
2222
else:
2323
return None
2424

25-
def get_timestep_embedding(timesteps, embedding_dim):
25+
def get_timestep_embedding(timesteps, embedding_dim, flip_sin_to_cos=False, downscale_freq_shift=1):
2626
"""
2727
This matches the implementation in Denoising Diffusion Probabilistic Models:
2828
From Fairseq.
@@ -33,11 +33,13 @@ def get_timestep_embedding(timesteps, embedding_dim):
3333
assert len(timesteps.shape) == 1
3434

3535
half_dim = embedding_dim // 2
36-
emb = math.log(10000) / (half_dim - 1)
36+
emb = math.log(10000) / (half_dim - downscale_freq_shift)
3737
emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb)
3838
emb = emb.to(device=timesteps.device)
3939
emb = timesteps.float()[:, None] * emb[None, :]
4040
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
41+
if flip_sin_to_cos:
42+
emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
4143
if embedding_dim % 2 == 1: # zero pad
4244
emb = torch.nn.functional.pad(emb, (0,1,0,0))
4345
return emb

‎comfy/ldm/seedvr/attention.py‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
import torch
2+
3+
from comfy.ldm.modules import attention as _attention
4+
5+
6+
def _var_attention_qkv(q, k, v, heads, skip_reshape):
7+
if skip_reshape:
8+
return q, k, v, q.shape[-1]
9+
total_tokens, embed_dim = q.shape
10+
head_dim = embed_dim // heads
11+
return (
12+
q.view(total_tokens, heads, head_dim),
13+
k.view(k.shape[0], heads, head_dim),
14+
v.view(v.shape[0], heads, head_dim),
15+
head_dim,
16+
)
17+
18+
19+
def _var_attention_output(out, heads, head_dim, skip_output_reshape):
20+
if skip_output_reshape:
21+
return out
22+
return out.reshape(-1, heads * head_dim)
23+
24+
25+
def var_attention_optimized_split(q, k, v, heads, cu_seqlens_q, cu_seqlens_k, *args, skip_reshape=False, skip_output_reshape=False, **kwargs):
26+
q, k, v, head_dim = _var_attention_qkv(q, k, v, heads, skip_reshape)
27+
28+
q_split_indices = cu_seqlens_q[1:-1]
29+
k_split_indices = cu_seqlens_k[1:-1]
30+
if k.shape[0] != v.shape[0]:
31+
raise ValueError("cu_seqlens_k does not match v token count")
32+
33+
q_splits = torch.tensor_split(q, q_split_indices, dim=0)
34+
k_splits = torch.tensor_split(k, k_split_indices, dim=0)
35+
v_splits = torch.tensor_split(v, k_split_indices, dim=0)
36+
if len(q_splits) != len(k_splits) or len(q_splits) != len(v_splits):
37+
raise ValueError("cu_seqlens_q and cu_seqlens_k must describe the same sequence count")
38+
39+
out = []
40+
for q_i, k_i, v_i in zip(q_splits, k_splits, v_splits):
41+
q_i = q_i.permute(1, 0, 2).unsqueeze(0)
42+
k_i = k_i.permute(1, 0, 2).unsqueeze(0)
43+
v_i = v_i.permute(1, 0, 2).unsqueeze(0)
44+
out_i = _attention.optimized_attention(q_i, k_i, v_i, heads, skip_reshape=True, skip_output_reshape=True)
45+
out.append(out_i.squeeze(0).permute(1, 0, 2))
46+
47+
out = torch.cat(out, dim=0)
48+
return _var_attention_output(out, heads, head_dim, skip_output_reshape)
49+
50+
51+
optimized_var_attention = var_attention_optimized_split

‎comfy/ldm/seedvr/color_fix.py‎

Lines changed: 301 additions & 0 deletions

0 commit comments

Comments
 (0)