Revert "Add SeedVR2 support (CORE-6) (Comfy-Org#14110)" (Comfy-Org#14359)

This reverts commit 7863cf0.

Author: comfyanonymous

comfy/latent_formats.py

@@ -4,7 +4,6 @@ class LatentFormat:
     scale_factor = 1.0
     latent_channels = 4
     latent_dimensions = 2
-    preserve_empty_channel_multiples = False
     latent_rgb_factors = None
     latent_rgb_factors_bias = None
     latent_rgb_factors_reshape = None
@@ -780,10 +779,6 @@ class ACEAudio(LatentFormat):
     latent_channels = 8
     latent_dimensions = 2
 
-class SeedVR2(LatentFormat):
-    latent_channels = 16
-    preserve_empty_channel_multiples = True
-
 class ACEAudio15(LatentFormat):
     latent_channels = 64
     latent_dimensions = 1

comfy/ldm/modules/attention.py

@@ -735,86 +735,7 @@ def attention_flash(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
         )
     return out
 
-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 _use_blackwell_attention():
-    device = model_management.get_torch_device()
-    if device.type != "cuda":
-        return False
-    major, minor = torch.cuda.get_device_capability(device)
-    return (major, minor) >= (12, 0)
-
-
-def _validate_split_cu_seqlens(name, cu_seqlens, token_count):
-    if cu_seqlens.dtype not in (torch.int32, torch.int64):
-        raise ValueError(f"{name} must use an integer dtype")
-    if cu_seqlens.ndim != 1 or cu_seqlens.numel() < 2:
-        raise ValueError(f"{name} must be a 1D tensor with at least two offsets")
-    if cu_seqlens[0].item() != 0:
-        raise ValueError(f"{name} must start at 0")
-    if (cu_seqlens[1:] <= cu_seqlens[:-1]).any().item():
-        raise ValueError(f"{name} must be strictly increasing")
-    if cu_seqlens[-1].item() != token_count:
-        raise ValueError(f"{name} does not match token count")
-
-
-def _split_indices(cu_seqlens):
-    return cu_seqlens[1:-1].to(device="cpu", dtype=torch.long)
-
-
-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)
-
-    _validate_split_cu_seqlens("cu_seqlens_q", cu_seqlens_q, q.shape[0])
-    _validate_split_cu_seqlens("cu_seqlens_k", cu_seqlens_k, k.shape[0])
-    if cu_seqlens_k[-1].item() != v.shape[0]:
-        raise ValueError("cu_seqlens_k does not match v token count")
-
-    q_split_indices = _split_indices(cu_seqlens_q)
-    k_split_indices = _split_indices(cu_seqlens_k)
-    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_dtype = q_i.dtype
-        if optimized_attention is attention_sage and q_i.dtype not in (torch.float16, torch.bfloat16):
-            q_i = q_i.to(torch.bfloat16)
-            k_i = k_i.to(torch.bfloat16)
-            v_i = v_i.to(torch.bfloat16)
-        out_i = optimized_attention(q_i, k_i, v_i, heads, skip_reshape=True, skip_output_reshape=True)
-        if out_i.dtype != out_dtype:
-            out_i = out_i.to(out_dtype)
-        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
 optimized_attention = attention_basic
 
 if model_management.sage_attention_enabled():
@@ -837,8 +758,6 @@ def var_attention_optimized_split(q, k, v, heads, cu_seqlens_q, cu_seqlens_k, *a
         logging.info("Using sub quadratic optimization for attention, if you have memory or speed issues try using: --use-split-cross-attention")
         optimized_attention = attention_sub_quad
 
-logging.info("Using optimized_attention split-loop for variable-length attention")
-
 optimized_attention_masked = optimized_attention
 
 
@@ -854,7 +773,6 @@ def var_attention_optimized_split(q, k, v, heads, cu_seqlens_q, cu_seqlens_k, *a
 register_attention_function("pytorch", attention_pytorch)
 register_attention_function("sub_quad", attention_sub_quad)
 register_attention_function("split", attention_split)
-register_attention_function("var_attention_optimized_split", var_attention_optimized_split)
 
 
 def optimized_attention_for_device(device, mask=False, small_input=False):
@@ -1291,3 +1209,5 @@ def forward(
             x = self.proj_out(x)
         out = x + x_in
         return out
+
+

comfy/ldm/modules/diffusionmodules/model.py

@@ -13,7 +13,6 @@
     import xformers
     import xformers.ops
 
-
 def torch_cat_if_needed(xl, dim):
     xl = [x for x in xl if x is not None and x.shape[dim] > 0]
     if len(xl) > 1:
@@ -23,8 +22,7 @@ def torch_cat_if_needed(xl, dim):
     else:
         return None
 
-
-def get_timestep_embedding(timesteps, embedding_dim, flip_sin_to_cos=False, downscale_freq_shift=1):
+def get_timestep_embedding(timesteps, embedding_dim):
     """
     This matches the implementation in Denoising Diffusion Probabilistic Models:
     From Fairseq.
@@ -35,13 +33,11 @@ def get_timestep_embedding(timesteps, embedding_dim, flip_sin_to_cos=False, down
     assert len(timesteps.shape) == 1
 
     half_dim = embedding_dim // 2
-    emb = math.log(10000) / (half_dim - downscale_freq_shift)
+    emb = math.log(10000) / (half_dim - 1)
     emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb)
     emb = emb.to(device=timesteps.device)
     emb = timesteps.float()[:, None] * emb[None, :]
     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
         emb = torch.nn.functional.pad(emb, (0,1,0,0))
     return emb