Spaces:
Runtime error
Runtime error
| import torch | |
| import torch.nn as nn | |
| from torchtyping import TensorType | |
| from einops import rearrange | |
| from .image_encoders import get_image_encoder | |
| from .config import MultimodalConfig | |
| # ------------------------- Image prefix ---------------------------------- | |
| # for models that are fixed to a specific sequence lengths (i.e clip models with no pooling), the sequence lengths are below | |
| ENCODER_SEQ_LENS = { | |
| "clip_resnet": 49, | |
| "clip_resnet_large": 144, | |
| } | |
| ENCODER_OUT_DIMS = { | |
| "nfresnet50": 2048, | |
| "clip": 512, | |
| "clip_resnet": 2560, | |
| "clip_resnet_large": 3072, | |
| } | |
| class ImagePrefix(nn.Module): | |
| """ | |
| Takes in a batch of images and returns a batch of embeddings of the | |
| same dimensions as the LM's word embeddings. | |
| :param config: MultimodalConfig object | |
| :param out_dim: output dimension of the embedding | |
| :param device: device to run the model on | |
| """ | |
| def __init__( | |
| self, | |
| config: MultimodalConfig, | |
| out_dim: int = 2048, | |
| device=None, | |
| ): | |
| super().__init__() | |
| self.device = device or torch.device( | |
| "cuda" if torch.cuda.is_available() else "cpu" | |
| ) | |
| self.config = config | |
| self.encoder_type = config.encoder_name | |
| # get image encoder backbone | |
| self.enc = get_image_encoder( | |
| config.encoder_name, | |
| pretrained=config.pretrained_img_encoder, | |
| ) | |
| self.encoder_out_dim = ENCODER_OUT_DIMS[ | |
| self.encoder_type | |
| ] # out dim for image encoder | |
| self.out_dim = out_dim # out dim for lm | |
| # set the out seq len to that specified in the config, or for some models, the hardcoded value | |
| self.out_seq_len = ( | |
| config.image_seq_len | |
| if config.encoder_name not in ENCODER_SEQ_LENS | |
| else ENCODER_SEQ_LENS[config.encoder_name] | |
| ) | |
| # get the output projection | |
| proj_out_dim = ( | |
| (self.out_dim * self.out_seq_len) | |
| if self.encoder_type not in ENCODER_SEQ_LENS | |
| else self.out_dim | |
| ) | |
| self.proj = nn.Linear(self.encoder_out_dim, proj_out_dim) | |
| self.dropout = nn.Dropout(config.image_embed_dropout_prob) | |
| self.use_layernorm = config.use_image_embed_layernorm | |
| if self.use_layernorm: | |
| self.ln = nn.LayerNorm(self.out_dim) | |
| def forward( | |
| self, x: TensorType["b", "c", "h", "w"] | |
| ) -> TensorType["b", "seq", "out_dim"]: | |
| # pass through image encoder | |
| logits = self.enc(x) | |
| # remove trailing dimensions of size 1 + pass through linear | |
| if logits.ndim == 4: | |
| logits = rearrange(logits, "b d 1 1 -> b d") | |
| elif logits.ndim == 3: | |
| assert self.encoder_type in ENCODER_SEQ_LENS | |
| else: | |
| assert logits.ndim == 2 | |
| logits = self.proj(logits) | |
| # reshape to desired output shape | |
| if ( | |
| self.encoder_type not in ENCODER_SEQ_LENS | |
| ): # don't need to reshape those with fixed seq lens / no pooling | |
| logits = rearrange( | |
| logits, "b (s d) -> b s d", d=self.out_dim, s=self.out_seq_len | |
| ) | |
| # pass through dropout and layer norm | |
| logits = self.dropout(logits) | |
| if self.use_layernorm: | |
| logits = self.ln(logits) | |
| return logits | |