Ismayil/optimize streams syncs

pyproject.toml

@@ -7,7 +7,7 @@ name = "seedvr"
 version = "1.0.6"
 description = "SeedVR: Seeding Infinity in Diffusion Transformer Towards Generic Video Restoration"
 readme = "readme.md"
-requires-python = ">=3.9,<3.11"
+requires-python = ">=3.9,<3.12"
 license = {text = "Apache-2.0"}
 authors = [
     {name = "Jianyi Wang", email = "iceclear@bytedance.com"},
@@ -38,6 +38,7 @@ classifiers = [
     "Programming Language :: Python :: 3",
     "Programming Language :: Python :: 3.9",
     "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
     "Topic :: Scientific/Engineering :: Artificial Intelligence",
     "Topic :: Scientific/Engineering :: Image Processing",
     "Topic :: Multimedia :: Video",
@@ -97,7 +98,10 @@ Documentation = "https://iceclear.github.io/projects/seedvr/"
 "ComfyUI Integration" = "https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler"
 
 [tool.setuptools]
-packages = ["seedvr"]
+include-package-data = true
+
+[tool.setuptools.packages.find]
+include = ["seedvr*"]
 
 [tool.setuptools.package-data]
 seedvr = [
@@ -147,7 +151,4 @@ addopts = [
     "--cov=seedvr",
     "--cov-report=term-missing",
     "--cov-report=html",
-]
-
-
-
+]

seedvr/models/dit/attention.py

@@ -14,21 +14,8 @@
 
 import torch
 import torch.nn.functional as F
-
-try:
-    from flash_attn_interface import flash_attn_varlen_func
-
-    print("Using FlashAttention3")
-except ImportError:
-    try:
-        from flash_attn import flash_attn_varlen_func
-
-        print("Using FlashAttention2")
-    except ImportError:
-        print("Using PyTorch Attention")
-        flash_attn_varlen_func = None
-
 from torch import nn
+from torch.nn.attention.varlen import varlen_attn
 
 
 class TorchAttention(nn.Module):
@@ -56,7 +43,7 @@ def forward(self, *args, **kwargs):
         )
 
 
-class FlashAttentionVarlen(nn.Module):
+class VarlenAttention(nn.Module):
     def tflops(self, args, kwargs, output) -> float:
         cu_seqlens_q = kwargs["cu_seqlens_q"]
         cu_seqlens_k = kwargs["cu_seqlens_k"]
@@ -66,7 +53,26 @@ def tflops(self, args, kwargs, output) -> float:
         return h * (4 * d * (seqlens_q * seqlens_k).sum())
 
     def forward(self, *args, **kwargs):
-        kwargs["deterministic"] = torch.are_deterministic_algorithms_enabled()
-        if flash_attn_varlen_func is None:
-            return TorchAttention()(*args, **kwargs)
-        return flash_attn_varlen_func(*args, **kwargs)
+        query = kwargs.pop("query", kwargs.pop("q", None))
+        key = kwargs.pop("key", kwargs.pop("k", None))
+        value = kwargs.pop("value", kwargs.pop("v", None))
+        if query is None and len(args) > 0:
+            query = args[0]
+        if key is None and len(args) > 1:
+            key = args[1]
+        if value is None and len(args) > 2:
+            value = args[2]
+        if query is None or key is None or value is None:
+            raise ValueError("query, key, value must be provided")
+
+        return varlen_attn(
+            query=query,
+            key=key,
+            value=value,
+            cu_seq_q=kwargs.pop("cu_seqlens_q"),
+            cu_seq_k=kwargs.pop("cu_seqlens_k"),
+            max_q=kwargs.pop("max_seqlen_q"),
+            max_k=kwargs.pop("max_seqlen_k"),
+            is_causal=kwargs.pop("is_causal", False),
+            return_aux=kwargs.pop("return_aux", None),
+        )

seedvr/models/dit/na.py

@@ -26,7 +26,12 @@ def flatten(
     torch.LongTensor,  # (b n)
 ]:
     assert len(hid) > 0
-    shape = torch.stack([torch.tensor(x.shape[:-1], device=hid[0].device) for x in hid])
+    shape = torch.stack(
+        [
+            torch.tensor(x.shape[:-1], device="cpu", pin_memory=True)
+            for x in hid
+        ]
+    ).pin_memory()
     hid = torch.cat([x.flatten(0, -2) for x in hid])
     return hid, shape
 
@@ -55,13 +60,17 @@ def concat(
 def concat_idx(
     vid_len: torch.LongTensor,  # (b)
     txt_len: torch.LongTensor,  # (b)
+    device: torch.device,
 ) -> tuple[
     Callable,
     Callable,
 ]:
-    device = vid_len.device
-    vid_idx = torch.arange(vid_len.sum(), device=device)
-    txt_idx = torch.arange(len(vid_idx), len(vid_idx) + txt_len.sum(), device=device)
+    vid_idx = torch.arange(int(vid_len.sum()), device=device)
+    txt_idx = torch.arange(
+        len(vid_idx),
+        len(vid_idx) + int(txt_len.sum()),
+        device=device,
+    )
     tgt_idx = concat(vid_idx, txt_idx, vid_len, txt_len)
     src_idx = torch.argsort(tgt_idx)
     return (
@@ -105,13 +114,17 @@ def repeat_concat_idx(
     vid_len: torch.LongTensor,  # (n*b)
     txt_len: torch.LongTensor,  # (b)
     txt_repeat: torch.LongTensor,  # (n)
+    device: torch.device,
 ) -> tuple[
     Callable,
     Callable,
 ]:
-    device = vid_len.device
-    vid_idx = torch.arange(vid_len.sum(), device=device)
-    txt_idx = torch.arange(len(vid_idx), len(vid_idx) + txt_len.sum(), device=device)
+    vid_idx = torch.arange(int(vid_len.sum()), device=device)
+    txt_idx = torch.arange(
+        len(vid_idx),
+        len(vid_idx) + int(txt_len.sum()),
+        device=device,
+    )
     txt_repeat_list = txt_repeat.tolist()
     tgt_idx = repeat_concat(vid_idx, txt_idx, vid_len, txt_len, txt_repeat)
     src_idx = torch.argsort(tgt_idx)
@@ -160,11 +173,10 @@ def rearrange(
 def rearrange_idx(
     hid_shape: torch.LongTensor,  # (b n)
     pattern: str,
+    device: torch.device,
     **kwargs: dict[str, int],
 ) -> tuple[Callable, Callable, torch.LongTensor]:
-    hid_idx = torch.arange(hid_shape.prod(-1).sum(), device=hid_shape.device).unsqueeze(
-        -1
-    )
+    hid_idx = torch.arange(int(hid_shape.prod(-1).sum()), device=device).unsqueeze(-1)
     tgt_idx, tgt_shape = rearrange(hid_idx, hid_shape, pattern, **kwargs)
     tgt_idx = tgt_idx.squeeze(-1)
     src_idx = torch.argsort(tgt_idx)
@@ -227,18 +239,17 @@ def window(
 ):
     hid = unflatten(hid, hid_shape)
     hid = list(map(window_fn, hid))
-    hid_windows = torch.tensor(list(map(len, hid)), device=hid_shape.device)
+    hid_windows = torch.tensor(list(map(len, hid)), device="cpu", pin_memory=True)
     hid, hid_shape = flatten(list(chain(*hid)))
     return hid, hid_shape, hid_windows
 
 
 def window_idx(
     hid_shape: torch.LongTensor,  # (b n)
     window_fn: Callable[[torch.Tensor], list[torch.Tensor]],
+    device: torch.device,
 ):
-    hid_idx = torch.arange(hid_shape.prod(-1).sum(), device=hid_shape.device).unsqueeze(
-        -1
-    )
+    hid_idx = torch.arange(int(hid_shape.prod(-1).sum()), device=device).unsqueeze(-1)
     tgt_idx, tgt_shape, tgt_windows = window(hid_idx, hid_shape, window_fn)
     tgt_idx = tgt_idx.squeeze(-1)
     src_idx = torch.argsort(tgt_idx)

seedvr/models/dit/nablocks/mmsr_block.py

@@ -25,7 +25,7 @@
 from seedvr.common.utils import get_torch_dtype, safe_pad_operation
 
 from .. import na
-from ..attention import FlashAttentionVarlen
+from ..attention import VarlenAttention
 from ..blocks.mmdit_window_block import MMWindowAttention, MMWindowTransformerBlock
 from ..mm import MMArg
 from ..modulation import ada_layer_type
@@ -65,7 +65,7 @@ def __init__(
             shared_qkv=shared_qkv,
         )
         self.rope = NaRotaryEmbedding3d(dim=head_dim // 2) if qk_rope else None
-        self.attn = FlashAttentionVarlen()
+        self.attn = VarlenAttention()
         self.window_op = get_window_op(window_method)
         self.rope.to(get_torch_dtype(dtype))
 
@@ -105,7 +105,7 @@ def make_window(x: torch.Tensor):
 
         window_partition, window_reverse, window_shape, window_count = cache_win(
             "win_transform",
-            lambda: na.window_idx(vid_shape, make_window),
+            lambda: na.window_idx(vid_shape, make_window, device=vid_qkv.device),
         )
         vid_qkv_win = window_partition(vid_qkv)
 
@@ -128,7 +128,13 @@ def make_window(x: torch.Tensor):
         )
         all_len_win = cache_win("all_len", lambda: vid_len_win + txt_len_win)
         concat_win, unconcat_win = cache_win(
-            "mm_pnp", lambda: na.repeat_concat_idx(vid_len_win, txt_len, window_count)
+            "mm_pnp",
+            lambda: na.repeat_concat_idx(
+                vid_len_win,
+                txt_len,
+                window_count,
+                device=vid_q.device,
+            ),
         )
 
         # window rope
@@ -141,11 +147,17 @@ def make_window(x: torch.Tensor):
             v=concat_win(vid_v, txt_v).bfloat16(),
             cu_seqlens_q=cache_win(
                 "vid_seqlens_q",
-                lambda: safe_pad_operation(all_len_win.cumsum(0), (1, 0)).int(),
+                lambda: safe_pad_operation(all_len_win.cumsum(0), (1, 0))
+                .int()
+                .pin_memory()
+                .to(device=vid_q.device, non_blocking=True),
             ),
             cu_seqlens_k=cache_win(
                 "vid_seqlens_k",
-                lambda: safe_pad_operation(all_len_win.cumsum(0), (1, 0)).int(),
+                lambda: safe_pad_operation(all_len_win.cumsum(0), (1, 0))
+                .int()
+                .pin_memory()
+                .to(device=vid_q.device, non_blocking=True),
             ),
             max_seqlen_q=cache_win(
                 "vid_max_seqlen_q", lambda: all_len_win.max().item()

seedvr/models/dit_v2/attention.py

@@ -14,8 +14,8 @@
 
 import torch
 import torch.nn.functional as F
-from flash_attn import flash_attn_varlen_func
 from torch import nn
+from torch.nn.attention.varlen import varlen_attn
 
 
 class TorchAttention(nn.Module):
@@ -33,7 +33,7 @@ def forward(self, *args, **kwargs):
         return F.scaled_dot_product_attention(*args, **kwargs)
 
 
-class FlashAttentionVarlen(nn.Module):
+class VarlenAttention(nn.Module):
     def tflops(self, args, kwargs, output) -> float:
         cu_seqlens_q = kwargs["cu_seqlens_q"]
         cu_seqlens_k = kwargs["cu_seqlens_k"]
@@ -43,5 +43,26 @@ def tflops(self, args, kwargs, output) -> float:
         return h * (4 * d * (seqlens_q * seqlens_k).sum())
 
     def forward(self, *args, **kwargs):
-        kwargs["deterministic"] = torch.are_deterministic_algorithms_enabled()
-        return flash_attn_varlen_func(*args, **kwargs)
+        query = kwargs.pop("query", kwargs.pop("q", None))
+        key = kwargs.pop("key", kwargs.pop("k", None))
+        value = kwargs.pop("value", kwargs.pop("v", None))
+        if query is None and len(args) > 0:
+            query = args[0]
+        if key is None and len(args) > 1:
+            key = args[1]
+        if value is None and len(args) > 2:
+            value = args[2]
+        if query is None or key is None or value is None:
+            raise ValueError("query, key, value must be provided")
+
+        return varlen_attn(
+            query=query,
+            key=key,
+            value=value,
+            cu_seq_q=kwargs.pop("cu_seqlens_q"),
+            cu_seq_k=kwargs.pop("cu_seqlens_k"),
+            max_q=kwargs.pop("max_seqlen_q"),
+            max_k=kwargs.pop("max_seqlen_k"),
+            is_causal=kwargs.pop("is_causal", False),
+            return_aux=kwargs.pop("return_aux", None),
+        )

seedvr/models/dit_v2/na.py

@@ -26,7 +26,12 @@ def flatten(
     torch.LongTensor,  # (b n)
 ]:
     assert len(hid) > 0
-    shape = torch.stack([torch.tensor(x.shape[:-1], device=hid[0].device) for x in hid])
+    shape = torch.stack(
+        [
+            torch.tensor(x.shape[:-1], device="cpu", pin_memory=True)
+            for x in hid
+        ]
+    ).pin_memory()
     hid = torch.cat([x.flatten(0, -2) for x in hid])
     return hid, shape
 
@@ -55,13 +60,17 @@ def concat(
 def concat_idx(
     vid_len: torch.LongTensor,  # (b)
     txt_len: torch.LongTensor,  # (b)
+    device: torch.device,
 ) -> tuple[
     Callable,
     Callable,
 ]:
-    device = vid_len.device
-    vid_idx = torch.arange(vid_len.sum(), device=device)
-    txt_idx = torch.arange(len(vid_idx), len(vid_idx) + txt_len.sum(), device=device)
+    vid_idx = torch.arange(int(vid_len.sum()), device=device)
+    txt_idx = torch.arange(
+        len(vid_idx),
+        len(vid_idx) + int(txt_len.sum()),
+        device=device,
+    )
     tgt_idx = concat(vid_idx, txt_idx, vid_len, txt_len)
     src_idx = torch.argsort(tgt_idx)
     return (
@@ -105,13 +114,17 @@ def repeat_concat_idx(
     vid_len: torch.LongTensor,  # (n*b)
     txt_len: torch.LongTensor,  # (b)
     txt_repeat: torch.LongTensor,  # (n)
+    device: torch.device,
 ) -> tuple[
     Callable,
     Callable,
 ]:
-    device = vid_len.device
-    vid_idx = torch.arange(vid_len.sum(), device=device)
-    txt_idx = torch.arange(len(vid_idx), len(vid_idx) + txt_len.sum(), device=device)
+    vid_idx = torch.arange(int(vid_len.sum()), device=device)
+    txt_idx = torch.arange(
+        len(vid_idx),
+        len(vid_idx) + int(txt_len.sum()),
+        device=device,
+    )
     txt_repeat_list = txt_repeat.tolist()
     tgt_idx = repeat_concat(vid_idx, txt_idx, vid_len, txt_len, txt_repeat)
     src_idx = torch.argsort(tgt_idx)
@@ -160,11 +173,10 @@ def rearrange(
 def rearrange_idx(
     hid_shape: torch.LongTensor,  # (b n)
     pattern: str,
+    device: torch.device,
     **kwargs: dict[str, int],
 ) -> tuple[Callable, Callable, torch.LongTensor]:
-    hid_idx = torch.arange(hid_shape.prod(-1).sum(), device=hid_shape.device).unsqueeze(
-        -1
-    )
+    hid_idx = torch.arange(int(hid_shape.prod(-1).sum()), device=device).unsqueeze(-1)
     tgt_idx, tgt_shape = rearrange(hid_idx, hid_shape, pattern, **kwargs)
     tgt_idx = tgt_idx.squeeze(-1)
     src_idx = torch.argsort(tgt_idx)
@@ -227,18 +239,17 @@ def window(
 ):
     hid = unflatten(hid, hid_shape)
     hid = list(map(window_fn, hid))
-    hid_windows = torch.tensor(list(map(len, hid)), device=hid_shape.device)
+    hid_windows = torch.tensor(list(map(len, hid)), device="cpu", pin_memory=True)
     hid, hid_shape = flatten(list(chain(*hid)))
     return hid, hid_shape, hid_windows
 
 
 def window_idx(
     hid_shape: torch.LongTensor,  # (b n)
     window_fn: Callable[[torch.Tensor], list[torch.Tensor]],
+    device: torch.device,
 ):
-    hid_idx = torch.arange(hid_shape.prod(-1).sum(), device=hid_shape.device).unsqueeze(
-        -1
-    )
+    hid_idx = torch.arange(int(hid_shape.prod(-1).sum()), device=device).unsqueeze(-1)
     tgt_idx, tgt_shape, tgt_windows = window(hid_idx, hid_shape, window_fn)
     tgt_idx = tgt_idx.squeeze(-1)
     src_idx = torch.argsort(tgt_idx)