Spaces:
Runtime error
Runtime error
| # Copyright (c) OpenMMLab. All rights reserved. | |
| """Taken from https://github.com/lucidrains/flamingo-pytorch.""" | |
| from typing import Optional | |
| import torch | |
| from einops import rearrange, repeat | |
| from torch import einsum, nn | |
| def FeedForward(dim, mult: int = 4): | |
| """Feedforward layers. | |
| Args: | |
| mult (int): Layer expansion muliplier. Defaults to 4. | |
| """ | |
| inner_dim = int(dim * mult) | |
| return nn.Sequential( | |
| nn.LayerNorm(dim), | |
| nn.Linear(dim, inner_dim, bias=False), | |
| nn.GELU(), | |
| nn.Linear(inner_dim, dim, bias=False), | |
| ) | |
| class PerceiverAttention(nn.Module): | |
| """Perceiver attetion layers. | |
| Args: | |
| dim (int): Input dimensions. | |
| dim_head (int): Number of dimension heads. Defaults to 64. | |
| heads (int): Number of heads. Defaults to 8. | |
| """ | |
| def __init__(self, *, dim: int, dim_head: int = 64, heads: int = 8): | |
| super().__init__() | |
| self.scale = dim_head**-0.5 | |
| self.heads = heads | |
| inner_dim = dim_head * heads | |
| self.norm_media = nn.LayerNorm(dim) | |
| self.norm_latents = nn.LayerNorm(dim) | |
| self.to_q = nn.Linear(dim, inner_dim, bias=False) | |
| self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False) | |
| self.to_out = nn.Linear(inner_dim, dim, bias=False) | |
| def forward(self, x: torch.Tensor, latents: torch.Tensor): | |
| """Forward function. | |
| Args: | |
| x (torch.Tensor): image features of shape (b, T, n1, D). | |
| latent (torch.Tensor): latent features of shape (b, T, n2, D). | |
| """ | |
| x = self.norm_media(x) | |
| latents = self.norm_latents(latents) | |
| h = self.heads | |
| q = self.to_q(latents) | |
| kv_input = torch.cat((x, latents), dim=-2) | |
| k, v = self.to_kv(kv_input).chunk(2, dim=-1) | |
| q = rearrange(q, 'b t n (h d) -> b h t n d', h=h) | |
| k = rearrange(k, 'b t n (h d) -> b h t n d', h=h) | |
| v = rearrange(v, 'b t n (h d) -> b h t n d', h=h) | |
| q = q * self.scale | |
| # attention | |
| sim = einsum('... i d, ... j d -> ... i j', q, k) | |
| sim = sim - sim.amax(dim=-1, keepdim=True).detach() | |
| attn = sim.softmax(dim=-1) | |
| out = einsum('... i j, ... j d -> ... i d', attn, v) | |
| out = rearrange(out, 'b h t n d -> b t n (h d)', h=h) | |
| return self.to_out(out) | |
| class PerceiverResampler(nn.Module): | |
| """Perceiver resampler layers. | |
| Args: | |
| dim (int): Input dimensions. | |
| depth (int): Depth of resampler. Defaults to 6. | |
| dim_head (int): Number of dimension heads. Defaults to 64. | |
| heads (int): Number of heads. Defaults to 8. | |
| num_latents (int): Number of latents. Defaults to 64. | |
| max_num_media (int, optional): Max number of media. | |
| Defaults to None. | |
| max_num_frames (int, optional): Max number of frames. | |
| Defaults to None. | |
| ff_mult (int): Feed forward multiplier. Defaults to 4. | |
| """ | |
| def __init__( | |
| self, | |
| *, | |
| dim: int, | |
| depth: int = 6, | |
| dim_head: int = 64, | |
| heads: int = 8, | |
| num_latents: int = 64, | |
| max_num_media: Optional[int] = None, | |
| max_num_frames: Optional[int] = None, | |
| ff_mult: int = 4, | |
| ): | |
| super().__init__() | |
| self.latents = nn.Parameter(torch.randn(num_latents, dim)) | |
| self.frame_embs = ( | |
| nn.Parameter(torch.randn(max_num_frames, dim)) | |
| if max_num_frames is not None else None) | |
| self.media_time_embs = ( | |
| nn.Parameter(torch.randn(max_num_media, 1, dim)) | |
| if max_num_media is not None else None) | |
| self.layers = nn.ModuleList([]) | |
| for _ in range(depth): | |
| self.layers.append( | |
| nn.ModuleList([ | |
| PerceiverAttention( | |
| dim=dim, dim_head=dim_head, heads=heads), | |
| FeedForward(dim=dim, mult=ff_mult), | |
| ])) | |
| self.norm = nn.LayerNorm(dim) | |
| def forward(self, x: torch.Tensor): | |
| """Forward function for perceiver sampler. | |
| Args: | |
| x (torch.Tensor): image features of shape (b, T, F, v, D) | |
| Returns: | |
| torch.Tensor: shape (b, T, n, D) where n is self.num_latents | |
| """ | |
| b, T, F, v = x.shape[:4] | |
| # frame and media time embeddings | |
| if self.frame_embs is not None: | |
| frame_embs = repeat( | |
| self.frame_embs[:F], 'F d -> b T F v d', b=b, T=T, v=v) | |
| x = x + frame_embs | |
| x = rearrange(x, 'b T F v d -> b T (F v) d' | |
| ) # flatten the frame and spatial dimensions | |
| if self.media_time_embs is not None: | |
| x = x + self.media_time_embs[:T] | |
| # blocks | |
| latents = repeat(self.latents, 'n d -> b T n d', b=b, T=T) | |
| for attn, ff in self.layers: | |
| latents = attn(x, latents) + latents | |
| latents = ff(latents) + latents | |
| return self.norm(latents) | |
| class MaskedCrossAttention(nn.Module): | |
| """Masked cross attention layers. | |
| Args: | |
| dim (int): Input text feature dimensions. | |
| dim_visual (int): Input visual feature dimensions. | |
| dim_head (int): Number of dimension heads. Defaults to 64. | |
| heads (int): Number of heads. Defaults to 8. | |
| only_attend_immediate_media (bool): Whether attend immediate media. | |
| Defaults to True. | |
| """ | |
| def __init__( | |
| self, | |
| *, | |
| dim: int, | |
| dim_visual: int, | |
| dim_head: int = 64, | |
| heads: int = 8, | |
| only_attend_immediate_media: bool = True, | |
| ): | |
| super().__init__() | |
| self.scale = dim_head**-0.5 | |
| self.heads = heads | |
| inner_dim = dim_head * heads | |
| self.norm = nn.LayerNorm(dim) | |
| self.to_q = nn.Linear(dim, inner_dim, bias=False) | |
| self.to_kv = nn.Linear(dim_visual, inner_dim * 2, bias=False) | |
| self.to_out = nn.Linear(inner_dim, dim, bias=False) | |
| # whether for text to only attend to immediate preceding image | |
| # or all previous images | |
| self.only_attend_immediate_media = only_attend_immediate_media | |
| def forward(self, | |
| x: torch.Tensor, | |
| media: torch.Tensor, | |
| media_locations: Optional[torch.Tensor] = None, | |
| attend_previous: bool = True): | |
| """Forward function for perceiver sampler. | |
| Args: | |
| x (torch.Tensor): text features of shape (B, T_txt, D_txt). | |
| media (torch.Tensor): image features of shape | |
| (B, T_img, n, D_img) where n is the dim of the latents. | |
| media_locations (torch.Tensor, optional): boolean mask identifying | |
| the media tokens in x of shape (B, T_txt). Defaults to None. | |
| attend_previous (bool): If false, ignores immediately preceding | |
| image and starts attending when following image. | |
| Defaults to True. | |
| """ | |
| _, T_img, n = media.shape[:3] | |
| h = self.heads | |
| x = self.norm(x) | |
| q = self.to_q(x) | |
| media = rearrange(media, 'b t n d -> b (t n) d') | |
| k, v = self.to_kv(media).chunk(2, dim=-1) | |
| q = rearrange(q, 'b n (h d) -> b h n d', h=h) | |
| k = rearrange(k, 'b n (h d) -> b h n d', h=h) | |
| v = rearrange(v, 'b n (h d) -> b h n d', h=h) | |
| q = q * self.scale | |
| sim = einsum('... i d, ... j d -> ... i j', q, k) | |
| if media_locations is not None: | |
| # at each boolean of True, increment the time counter | |
| # (relative to media time) | |
| text_time = media_locations.cumsum(dim=-1) | |
| media_time = torch.arange(T_img, device=x.device) + 1 | |
| if not attend_previous: | |
| text_time[~media_locations] += 1 | |
| # make sure max is still the number of images in the sequence | |
| text_time[text_time > repeat( | |
| torch.count_nonzero(media_locations, dim=1), | |
| 'b -> b i', | |
| i=text_time.shape[1], | |
| )] = 0 | |
| # text time must equal media time if only attending to most | |
| # immediate image otherwise, as long as text time is greater than | |
| # media time (if attending to all previous images / media) | |
| mask_op = torch.eq if self.only_attend_immediate_media else torch.ge # noqa | |
| text_to_media_mask = mask_op( | |
| rearrange(text_time, 'b i -> b 1 i 1'), | |
| repeat(media_time, 'j -> 1 1 1 (j n)', n=n), | |
| ) | |
| sim = sim.masked_fill(~text_to_media_mask, | |
| -torch.finfo(sim.dtype).max) | |
| sim = sim - sim.amax(dim=-1, keepdim=True).detach() | |
| attn = sim.softmax(dim=-1) | |
| if media_locations is not None and self.only_attend_immediate_media: | |
| # any text without a preceding media needs to have | |
| # attention zeroed out | |
| text_without_media_mask = text_time == 0 | |
| text_without_media_mask = rearrange(text_without_media_mask, | |
| 'b i -> b 1 i 1') | |
| attn = attn.masked_fill(text_without_media_mask, 0.0) | |
| out = einsum('... i j, ... j d -> ... i d', attn, v) | |
| out = rearrange(out, 'b h n d -> b n (h d)') | |
| return self.to_out(out) | |
| class GatedCrossAttentionBlock(nn.Module): | |
| """Gated cross attention layers. | |
| Args: | |
| dim (int): Input text feature dimensions. | |
| dim_visual (int): Input visual feature dimensions. | |
| dim_head (int): Number of dimension heads. Defaults to 64. | |
| heads (int): Number of heads. Defaults to 8. | |
| ff_mult (int): Feed forward multiplier. Defaults to 4. | |
| only_attend_immediate_media (bool): Whether attend immediate media. | |
| Defaults to True. | |
| """ | |
| def __init__( | |
| self, | |
| *, | |
| dim: int, | |
| dim_visual: int, | |
| dim_head: int = 64, | |
| heads: int = 8, | |
| ff_mult: int = 4, | |
| only_attend_immediate_media: bool = True, | |
| ): | |
| super().__init__() | |
| self.attn = MaskedCrossAttention( | |
| dim=dim, | |
| dim_visual=dim_visual, | |
| dim_head=dim_head, | |
| heads=heads, | |
| only_attend_immediate_media=only_attend_immediate_media, | |
| ) | |
| self.attn_gate = nn.Parameter(torch.tensor([0.0])) | |
| self.ff = FeedForward(dim, mult=ff_mult) | |
| self.ff_gate = nn.Parameter(torch.tensor([0.0])) | |
| def forward(self, | |
| x: torch.Tensor, | |
| media: torch.Tensor, | |
| media_locations: Optional[torch.Tensor] = None, | |
| attend_previous: bool = True): | |
| """Forward function for perceiver sampler. | |
| Args: | |
| x (torch.Tensor): text features of shape (B, T_txt, D_txt). | |
| media (torch.Tensor): image features of shape | |
| (B, T_img, n, D_img) where n is the dim of the latents. | |
| media_locations (torch.Tensor, optional): boolean mask identifying | |
| the media tokens in x of shape (B, T_txt). Defaults to None. | |
| attend_previous (bool): If false, ignores immediately preceding | |
| image and starts attending when following image. | |
| Defaults to True. | |
| """ | |
| x = ( | |
| self.attn( | |
| x, | |
| media, | |
| media_locations=media_locations, | |
| attend_previous=attend_previous, | |
| ) * self.attn_gate.tanh() + x) | |
| x = self.ff(x) * self.ff_gate.tanh() + x | |
| return x | |
| class FlamingoLayer(nn.Module): | |
| """Faminogo layers. | |
| Args: | |
| gated_cross_attn_layer (nn.Module): Gated cross attention layer. | |
| decoder_layer (nn.Module): Decoder layer. | |
| """ | |
| def __init__(self, gated_cross_attn_layer: nn.Module, | |
| decoder_layer: nn.Module): | |
| super().__init__() | |
| self.gated_cross_attn_layer = gated_cross_attn_layer | |
| self.decoder_layer = decoder_layer | |
| self.vis_x = None | |
| self.media_locations = None | |
| def is_conditioned(self) -> bool: | |
| """Check whether the layer is conditioned.""" | |
| return self.vis_x is not None | |
| def condition_vis_x(self, vis_x): | |
| """Set condition vision features.""" | |
| self.vis_x = vis_x | |
| def condition_media_locations(self, media_locations): | |
| """Set condition media locations.""" | |
| self.media_locations = media_locations | |
| def condition_attend_previous(self, attend_previous): | |
| """Set attend previous.""" | |
| self.attend_previous = attend_previous | |
| def forward( | |
| self, | |
| lang_x: torch.Tensor, | |
| attention_mask: Optional[torch.Tensor] = None, | |
| **decoder_layer_kwargs, | |
| ): | |
| """Forward function. | |
| Args: | |
| lang_x (torch.Tensor): language inputs. | |
| attention_mask (torch.Tensor, optional): text attention mask. | |
| Defaults to None. | |
| **decoder_layer_kwargs: Other decoder layer keyword arguments. | |
| """ | |
| if self.gated_cross_attn_layer is None: | |
| return self.decoder_layer( | |
| lang_x, attention_mask=attention_mask, **decoder_layer_kwargs) | |
| if self.vis_x is None: | |
| raise ValueError('vis_x must be conditioned before forward pass') | |
| if self.media_locations is None: | |
| raise ValueError( | |
| 'media_locations must be conditioned before forward pass') | |
| lang_x = self.gated_cross_attn_layer( | |
| lang_x, | |
| self.vis_x, | |
| media_locations=self.media_locations, | |
| attend_previous=self.attend_previous, | |
| ) | |
| lang_x = self.decoder_layer( | |
| lang_x, attention_mask=attention_mask, **decoder_layer_kwargs) | |
| return lang_x | |