diff --git a/VAE/models/temporal_ae.py b/VAE/models/temporal_ae.py index 690ca94..5f52549 100644 --- a/VAE/models/temporal_ae.py +++ b/VAE/models/temporal_ae.py @@ -162,11 +162,9 @@ class VideoDecoder(nn.Module): )) block_in = block_out if curr_res in attn_resolutions: - attn.append(make_time_attn( + attn.append(make_attn( block_in, attn_type=attn_type, - alpha=self.alpha, - merge_strategy=self.merge_strategy, )) up = nn.Module() up.block = block @@ -217,7 +215,7 @@ class VideoDecoder(nn.Module): for i_block in range(self.num_res_blocks+1): h = self.up[i_level].block[i_block](h, temb, **kwargs) if len(self.up[i_level].attn) > 0: - h = self.up[i_level].attn[i_block](h, **kwargs) + h = self.up[i_level].attn[i_block](h) if i_level != 0: h = self.up[i_level].upsample(h) @@ -232,318 +230,6 @@ class VideoDecoder(nn.Module): h = torch.tanh(h) return h -class CrossAttention(nn.Module): - def __init__( - self, - query_dim, - context_dim=None, - heads=8, - dim_head=64, - dropout=0.0, - backend=None, - ): - super().__init__() - inner_dim = dim_head * heads - context_dim = context_dim or query_dim - - self.scale = dim_head**-0.5 - self.heads = heads - - self.to_q = nn.Linear(query_dim, inner_dim, bias=False) - self.to_k = nn.Linear(context_dim, inner_dim, bias=False) - self.to_v = nn.Linear(context_dim, inner_dim, bias=False) - - self.to_out = nn.Sequential( - nn.Linear(inner_dim, query_dim), nn.Dropout(dropout) - ) - self.backend = backend - - def forward( - self, - x, - context=None, - mask=None, - additional_tokens=None, - n_times_crossframe_attn_in_self=0, - ): - h = self.heads - - if additional_tokens is not None: - # get the number of masked tokens at the beginning of the output sequence - n_tokens_to_mask = additional_tokens.shape[1] - # add additional token - x = torch.cat([additional_tokens, x], dim=1) - - q = self.to_q(x) - context = context or x - k = self.to_k(context) - v = self.to_v(context) - - if n_times_crossframe_attn_in_self: - # reprogramming cross-frame attention as in https://arxiv.org/abs/2303.13439 - assert x.shape[0] % n_times_crossframe_attn_in_self == 0 - n_cp = x.shape[0] // n_times_crossframe_attn_in_self - k = repeat( - k[::n_times_crossframe_attn_in_self], "b ... -> (b n) ...", n=n_cp - ) - v = repeat( - v[::n_times_crossframe_attn_in_self], "b ... -> (b n) ...", n=n_cp - ) - - q, k, v = map(lambda t: rearrange(t, "b n (h d) -> b h n d", h=h), (q, k, v)) - - ## old - """ - sim = einsum('b i d, b j d -> b i j', q, k) * self.scale - del q, k - - if exists(mask): - mask = rearrange(mask, 'b ... -> b (...)') - max_neg_value = -torch.finfo(sim.dtype).max - mask = repeat(mask, 'b j -> (b h) () j', h=h) - sim.masked_fill_(~mask, max_neg_value) - - # some quote about attention, that I just about had enough of - sim = sim.softmax(dim=-1) - - out = einsum('b i j, b j d -> b i d', sim, v) - """ - ## new - with sdp_kernel(**BACKEND_MAP[self.backend]): - # print("dispatching into backend", self.backend, "q/k/v shape: ", q.shape, k.shape, v.shape) - out = F.scaled_dot_product_attention( - q, k, v, attn_mask=mask - ) # scale is dim_head ** -0.5 per default - - del q, k, v - out = rearrange(out, "b h n d -> b n (h d)", h=h) - - if additional_tokens is not None: - # remove additional token - out = out[:, n_tokens_to_mask:] - return self.to_out(out) - -class VideoBlock(AttnBlock): - def __init__( - self, in_channels: int, alpha: float = 0, merge_strategy: str = "learned" - ): - super().__init__(in_channels) - # no context, single headed, as in base class - self.time_mix_block = VideoTransformerBlock( - dim=in_channels, - n_heads=1, - d_head=in_channels, - checkpoint=False, - ff_in=True, - attn_mode="softmax", - ) - - time_embed_dim = self.in_channels * 4 - self.video_time_embed = torch.nn.Sequential( - torch.nn.Linear(self.in_channels, time_embed_dim), - torch.nn.SiLU(), - torch.nn.Linear(time_embed_dim, self.in_channels), - ) - - self.merge_strategy = merge_strategy - if self.merge_strategy == "fixed": - self.register_buffer("mix_factor", torch.Tensor([alpha])) - elif self.merge_strategy == "learned": - self.register_parameter( - "mix_factor", torch.nn.Parameter(torch.Tensor([alpha])) - ) - else: - raise ValueError(f"unknown merge strategy {self.merge_strategy}") - - def forward(self, x, timesteps, skip_video=False): - if skip_video: - return super().forward(x) - - x_in = x - x = self.attention(x) - h, w = x.shape[2:] - x = rearrange(x, "b c h w -> b (h w) c") - - x_mix = x - num_frames = torch.arange(timesteps, device=x.device) - num_frames = repeat(num_frames, "t -> b t", b=x.shape[0] // timesteps) - num_frames = rearrange(num_frames, "b t -> (b t)") - t_emb = timestep_embedding(num_frames, self.in_channels, repeat_only=False) - emb = self.video_time_embed(t_emb) # b, n_channels - emb = emb[:, None, :] - x_mix = x_mix + emb - - alpha = self.get_alpha() - x_mix = self.time_mix_block(x_mix, timesteps=timesteps) - x = alpha * x + (1.0 - alpha) * x_mix # alpha merge - - x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w) - x = self.proj_out(x) - - return x_in + x - - def attention(self, h_: torch.Tensor) -> torch.Tensor: - h_ = self.norm(h_) - q = self.q(h_) - k = self.k(h_) - v = self.v(h_) - - b, c, h, w = q.shape - q, k, v = map( - lambda x: rearrange(x, "b c h w -> b 1 (h w) c").contiguous(), (q, k, v) - ) - h_ = torch.nn.functional.scaled_dot_product_attention( - q, k, v - ) # scale is dim ** -0.5 per default - # compute attention - - return rearrange(h_, "b 1 (h w) c -> b c h w", h=h, w=w, c=c, b=b) - - def forward(self, x, **kwargs): - h_ = x - h_ = self.attention(h_) - h_ = self.proj_out(h_) - return x + h_ - - def get_alpha( - self, - ): - if self.merge_strategy == "fixed": - return self.mix_factor - elif self.merge_strategy == "learned": - return torch.sigmoid(self.mix_factor) - else: - raise NotImplementedError(f"unknown merge strategy {self.merge_strategy}") - -class VideoTransformerBlock(nn.Module): - ATTENTION_MODES = { - "softmax": CrossAttention, - # "softmax-xformers": MemoryEfficientCrossAttention, - } - - def __init__( - self, - dim, - n_heads, - d_head, - dropout=0.0, - context_dim=None, - gated_ff=True, - checkpoint=True, - timesteps=None, - ff_in=False, - inner_dim=None, - attn_mode="softmax", - disable_self_attn=False, - disable_temporal_crossattention=False, - switch_temporal_ca_to_sa=False, - ): - super().__init__() - - attn_cls = self.ATTENTION_MODES.get(attn_mode, "softmax") - - self.ff_in = ff_in or inner_dim is not None - if inner_dim is None: - inner_dim = dim - - assert int(n_heads * d_head) == inner_dim - - self.is_res = inner_dim == dim - - if self.ff_in: - self.norm_in = nn.LayerNorm(dim) - self.ff_in = FeedForward( - dim, dim_out=inner_dim, dropout=dropout, glu=gated_ff - ) - - self.timesteps = timesteps - self.disable_self_attn = disable_self_attn - if self.disable_self_attn: - self.attn1 = attn_cls( - query_dim=inner_dim, - heads=n_heads, - dim_head=d_head, - context_dim=context_dim, - dropout=dropout, - ) # is a cross-attention - else: - self.attn1 = attn_cls( - query_dim=inner_dim, heads=n_heads, dim_head=d_head, dropout=dropout - ) # is a self-attention - - self.ff = FeedForward(inner_dim, dim_out=dim, dropout=dropout, glu=gated_ff) - - if disable_temporal_crossattention: - if switch_temporal_ca_to_sa: - raise ValueError - else: - self.attn2 = None - else: - self.norm2 = nn.LayerNorm(inner_dim) - if switch_temporal_ca_to_sa: - self.attn2 = attn_cls( - query_dim=inner_dim, heads=n_heads, dim_head=d_head, dropout=dropout - ) # is a self-attention - else: - self.attn2 = attn_cls( - query_dim=inner_dim, - context_dim=context_dim, - heads=n_heads, - dim_head=d_head, - dropout=dropout, - ) # is self-attn if context is none - - self.norm1 = nn.LayerNorm(inner_dim) - self.norm3 = nn.LayerNorm(inner_dim) - self.switch_temporal_ca_to_sa = switch_temporal_ca_to_sa - - self.checkpoint = checkpoint - if self.checkpoint: - print(f"{self.__class__.__name__} is using checkpointing") - - def forward( - self, x: torch.Tensor, context: torch.Tensor = None, timesteps: int = None - ) -> torch.Tensor: - if self.checkpoint: - return checkpoint(self._forward, x, context, timesteps) - else: - return self._forward(x, context, timesteps=timesteps) - - def _forward(self, x, context=None, timesteps=None): - assert self.timesteps or timesteps - assert not (self.timesteps and timesteps) or self.timesteps == timesteps - timesteps = self.timesteps or timesteps - B, S, C = x.shape - x = rearrange(x, "(b t) s c -> (b s) t c", t=timesteps) - - if self.ff_in: - x_skip = x - x = self.ff_in(self.norm_in(x)) - if self.is_res: - x += x_skip - - if self.disable_self_attn: - x = self.attn1(self.norm1(x), context=context) + x - else: - x = self.attn1(self.norm1(x)) + x - - if self.attn2 is not None: - if self.switch_temporal_ca_to_sa: - x = self.attn2(self.norm2(x)) + x - else: - x = self.attn2(self.norm2(x), context=context) + x - x_skip = x - x = self.ff(self.norm3(x)) - if self.is_res: - x += x_skip - - x = rearrange( - x, "(b s) t c -> (b t) s c", s=S, b=B // timesteps, c=C, t=timesteps - ) - return x - - def get_last_layer(self): - return self.ff.net[-1].weight class ResBlock(nn.Module): """ @@ -781,44 +467,6 @@ class AE3DConv(torch.nn.Conv2d): x = self.time_mix_conv(x) return rearrange(x, "b c t h w -> (b t) c h w") -def make_time_attn(in_channels, attn_type="vanilla", attn_kwargs=None, alpha: float = 0, merge_strategy: str = "learned"): - if attn_type == "vanilla": - assert attn_kwargs is None - return VideoBlock( - in_channels, - alpha=alpha, - merge_strategy=merge_strategy, - ) - # lazy to add the xformers code - else: - return NotImplementedError() - -def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False): - """ - Create sinusoidal timestep embeddings. - :param timesteps: a 1-D Tensor of N indices, one per batch element. - These may be fractional. - :param dim: the dimension of the output. - :param max_period: controls the minimum frequency of the embeddings. - :return: an [N x dim] Tensor of positional embeddings. - """ - if not repeat_only: - half = dim // 2 - freqs = torch.exp( - -math.log(max_period) - * torch.arange(start=0, end=half, dtype=torch.float32) - / half - ).to(device=timesteps.device) - args = timesteps[:, None].float() * freqs[None] - embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) - if dim % 2: - embedding = torch.cat( - [embedding, torch.zeros_like(embedding[:, :1])], dim=-1 - ) - else: - embedding = repeat(timesteps, "b -> b d", d=dim) - return embedding - def normalization(channels): """ Make a standard normalization layer. @@ -854,31 +502,3 @@ def zero_module(module): for p in module.parameters(): p.detach().zero_() return module - -# feedforward -class GEGLU(nn.Module): - def __init__(self, dim_in, dim_out): - super().__init__() - self.proj = nn.Linear(dim_in, dim_out * 2) - - def forward(self, x): - x, gate = self.proj(x).chunk(2, dim=-1) - return x * nn.functional.gelu(gate) - -class FeedForward(nn.Module): - def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.0): - super().__init__() - inner_dim = int(dim * mult) - dim_out = dim_out or dim - project_in = ( - nn.Sequential(nn.Linear(dim, inner_dim), nn.GELU()) - if not glu - else GEGLU(dim, inner_dim) - ) - - self.net = nn.Sequential( - project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out) - ) - - def forward(self, x): - return self.net(x)