diff --git a/pyproject.toml b/pyproject.toml index b642233..dc294a4 100644 --- a/pyproject.toml +++ b/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", -] - - - +] \ No newline at end of file diff --git a/seedvr/models/dit/attention.py b/seedvr/models/dit/attention.py index 3699453..94b82ce 100644 --- a/seedvr/models/dit/attention.py +++ b/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), + ) diff --git a/seedvr/models/dit/na.py b/seedvr/models/dit/na.py index 5541dd5..db7feb6 100644 --- a/seedvr/models/dit/na.py +++ b/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,7 +239,7 @@ 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 @@ -235,10 +247,9 @@ def window( 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) diff --git a/seedvr/models/dit/nablocks/mmsr_block.py b/seedvr/models/dit/nablocks/mmsr_block.py index 1bd6906..06dcf14 100644 --- a/seedvr/models/dit/nablocks/mmsr_block.py +++ b/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() diff --git a/seedvr/models/dit_v2/attention.py b/seedvr/models/dit_v2/attention.py index abe3c1a..78c4434 100644 --- a/seedvr/models/dit_v2/attention.py +++ b/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), + ) diff --git a/seedvr/models/dit_v2/na.py b/seedvr/models/dit_v2/na.py index 5541dd5..db7feb6 100644 --- a/seedvr/models/dit_v2/na.py +++ b/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,7 +239,7 @@ 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 @@ -235,10 +247,9 @@ def window( 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) diff --git a/seedvr/models/dit_v2/nablocks/attention/mmattn.py b/seedvr/models/dit_v2/nablocks/attention/mmattn.py index 5709441..d6bd9c3 100644 --- a/seedvr/models/dit_v2/nablocks/attention/mmattn.py +++ b/seedvr/models/dit_v2/nablocks/attention/mmattn.py @@ -23,7 +23,7 @@ from seedvr.common.distributed.ops import gather_heads_scatter_seq, gather_seq_scatter_heads_qkv from seedvr.common.utils import safe_pad_operation from ... import na -from ...attention import FlashAttentionVarlen +from ...attention import VarlenAttention from ...mm import MMArg, MMModule from ...normalization import norm_layer_type from ...rope import get_na_rope @@ -71,7 +71,7 @@ def __init__( ) self.rope = get_na_rope(rope_type=rope_type, dim=rope_dim) - self.attn = FlashAttentionVarlen() + self.attn = VarlenAttention() def forward( self, @@ -118,14 +118,29 @@ def forward( txt_len = cache("txt_len", lambda: txt_shape.prod(-1)) all_len = cache("all_len", lambda: vid_len + txt_len) - concat, unconcat = cache("mm_pnp", lambda: na.concat_idx(vid_len, txt_len)) + concat, unconcat = cache( + "mm_pnp", + lambda: na.concat_idx(vid_len, txt_len, device=vid_q.device), + ) attn = self.attn( q=concat(vid_q, txt_q).bfloat16(), k=concat(vid_k, txt_k).bfloat16(), v=concat(vid_v, txt_v).bfloat16(), - cu_seqlens_q=cache("mm_seqlens", lambda: safe_pad_operation(all_len.cumsum(0), (1, 0)).int()), - cu_seqlens_k=cache("mm_seqlens", lambda: safe_pad_operation(all_len.cumsum(0), (1, 0)).int()), + cu_seqlens_q=cache( + "mm_seqlens", + lambda: safe_pad_operation(all_len.cumsum(0), (1, 0)) + .int() + .pin_memory() + .to(device=vid_q.device, non_blocking=True), + ), + cu_seqlens_k=cache( + "mm_seqlens", + lambda: safe_pad_operation(all_len.cumsum(0), (1, 0)) + .int() + .pin_memory() + .to(device=vid_q.device, non_blocking=True), + ), max_seqlen_q=cache("mm_maxlen", lambda: all_len.max().item()), max_seqlen_k=cache("mm_maxlen", lambda: all_len.max().item()), ).type_as(vid_q) @@ -190,7 +205,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) @@ -209,7 +224,13 @@ def make_window(x: torch.Tensor): txt_len_win = cache_win("txt_len", lambda: txt_len.repeat_interleave(window_count)) 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 @@ -217,16 +238,17 @@ def make_window(x: torch.Tensor): if self.rope.mm: # repeat text q and k for window mmrope _, num_h, _ = txt_q.shape + window_count_list = window_count.tolist() txt_q_repeat = rearrange(txt_q, "l h d -> l (h d)") txt_q_repeat = na.unflatten(txt_q_repeat, txt_shape) - txt_q_repeat = [[x] * n for x, n in zip(txt_q_repeat, window_count)] + txt_q_repeat = [[x] * n for x, n in zip(txt_q_repeat, window_count_list)] txt_q_repeat = list(chain(*txt_q_repeat)) txt_q_repeat, txt_shape_repeat = na.flatten(txt_q_repeat) txt_q_repeat = rearrange(txt_q_repeat, "l (h d) -> l h d", h=num_h) txt_k_repeat = rearrange(txt_k, "l h d -> l (h d)") txt_k_repeat = na.unflatten(txt_k_repeat, txt_shape) - txt_k_repeat = [[x] * n for x, n in zip(txt_k_repeat, window_count)] + txt_k_repeat = [[x] * n for x, n in zip(txt_k_repeat, window_count_list)] txt_k_repeat = list(chain(*txt_k_repeat)) txt_k_repeat, _ = na.flatten(txt_k_repeat) txt_k_repeat = rearrange(txt_k_repeat, "l (h d) -> l h d", h=num_h) @@ -242,10 +264,18 @@ def make_window(x: torch.Tensor): k=concat_win(vid_k, txt_k).bfloat16(), 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() + "vid_seqlens_q", + 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() + "vid_seqlens_k", + 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()), max_seqlen_k=cache_win("vid_max_seqlen_k", lambda: all_len_win.max().item()), diff --git a/seedvr/pipeline.py b/seedvr/pipeline.py index 59ee187..9a29159 100644 --- a/seedvr/pipeline.py +++ b/seedvr/pipeline.py @@ -144,9 +144,15 @@ def vae_encode( shift = getattr(self.vae.config, "shifting_factor", 0.0) if isinstance(scale, ListConfig): - scale = torch.tensor(scale, device=device, dtype=dtype) + scale = torch.tensor(scale, device="cpu", dtype=dtype, pin_memory=True).to( + device=device, + non_blocking=True, + ) if isinstance(shift, ListConfig): - shift = torch.tensor(shift, device=device, dtype=dtype) + shift = torch.tensor(shift, device="cpu", dtype=dtype, pin_memory=True).to( + device=device, + non_blocking=True, + ) # Group samples of the same shape to batches if enabled. if self.vae.grouping: @@ -210,9 +216,15 @@ def vae_decode( shift = getattr(self.vae.config, "shifting_factor", 0.0) if isinstance(scale, ListConfig): - scale = torch.tensor(scale, device=device, dtype=dtype) + scale = torch.tensor(scale, device="cpu", dtype=dtype, pin_memory=True).to( + device=device, + non_blocking=True, + ) if isinstance(shift, ListConfig): - shift = torch.tensor(shift, device=device, dtype=dtype) + shift = torch.tensor(shift, device="cpu", dtype=dtype, pin_memory=True).to( + device=device, + non_blocking=True, + ) # Group latents of the same shape to batches if enabled. if self.vae.grouping: @@ -271,10 +283,6 @@ def inference( latents, latents_shapes = flatten(noises) latents_cond, _ = flatten(conditions) - # Enter eval mode. - was_training = self.dit.training - self.dit.eval() - # Sampling. latents = self.sampler.sample( x=latents, @@ -298,9 +306,6 @@ def inference( ), ) - # Exit eval mode. - self.dit.train(was_training) - # Unflatten. latents = unflatten(latents, latents_shapes) latents = [latent.to(self.vae.dtype) for latent in latents] @@ -374,8 +379,16 @@ def add_noise( """ Add noise to the input. """ - t = torch.tensor([self.sampler.timesteps.T], device=x.device) * cond_noise_scale - shape = torch.tensor(x.shape[1:], device=x.device).unsqueeze(0) + t = ( + torch.tensor([self.sampler.timesteps.T], device="cpu", pin_memory=True) + .to(device=x.device, non_blocking=True) + * cond_noise_scale + ) + shape = ( + torch.tensor(x.shape[1:], device="cpu", pin_memory=True) + .to(device=x.device, non_blocking=True) + .unsqueeze(0) + ) t = self.timestep_transform(t, shape) x = self.sampler.schedule.forward(x, aug_noise, t) return x @@ -503,11 +516,15 @@ def __call__( samples = [sample.unsqueeze(1) if sample.ndim == 3 else sample for sample in samples] batch_media = rearrange(batch_media, "c t h w -> t c h w").to( - self.wavelet_kernel.device, dtype=self.wavelet_kernel.dtype + self.wavelet_kernel.device, + dtype=self.wavelet_kernel.dtype, + non_blocking=True, ) samples = [ rearrange(sample, "c t h w -> t c h w").to( - self.wavelet_kernel.device, dtype=self.wavelet_kernel.dtype + self.wavelet_kernel.device, + dtype=self.wavelet_kernel.dtype, + non_blocking=True, ) for sample in samples ]