fromage/models.py

from typing import Callable, List, Optional, Tuple, Union
from collections import namedtuple
import json
import glob
import math
import numpy as np
import os
import torch
from torch import Tensor
import torch.nn as nn
import torch.nn.functional as F
from einops import rearrange
from functools import partial
import pickle as pkl
from PIL import Image, UnidentifiedImageError

from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
from transformers import OPTForCausalLM, GPT2Tokenizer
from transformers import CLIPVisionModel, CLIPVisionConfig

from fromage import utils


class FrozenArgs:
  freeze_lm: bool = True
  freeze_vm: bool = True
  opt_version: str = 'facebook/opt-6.7b'
  visual_encoder: str = 'openai/clip-vit-large-patch14'
  n_visual_tokens: int = 1
  image_embed_dropout_prob: float = 0.0
  task: str = 'captioning'
  shared_emb_dim: Optional[int] = 256
  text_emb_layers: List[int] = [-1]
  retrieval_token_idx: int = 0


class FromageModel(nn.Module):
  def __init__(self, tokenizer, args: FrozenArgs = FrozenArgs()):
    super().__init__()
    self.tokenizer = tokenizer
    self.feature_extractor = utils.get_feature_extractor_for_model(args.visual_encoder, train=False)
    self.image_token = self.tokenizer.cls_token_id
    assert args.text_emb_layers != set(args.text_emb_layers), 'text_emb_layers not unique'
    self.args = args

    opt_version = args.opt_version
    visual_encoder = args.visual_encoder
    n_visual_tokens = args.n_visual_tokens
    print(f"Using {opt_version} for the language model.")
    print(f"Using {visual_encoder} for the visual model with {n_visual_tokens} visual tokens.")

    self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / 0.07))

    if 'facebook/opt' in opt_version:
      self.lm = OPTForCausalLM.from_pretrained(opt_version)
    else:
      raise NotImplementedError

    self.opt_version = opt_version

    if self.args.freeze_lm:
      self.lm.eval()
      print("Freezing the LM.")
      for param in self.lm.parameters():
        param.requires_grad = False
    else:
      self.lm.train()

    self.retrieval_token_idx = args.retrieval_token_idx
    print(f'Initializing embedding for the retrieval token [RET] (id = {self.retrieval_token_idx}).')
    self.lm.resize_token_embeddings(len(tokenizer))

    self.input_embeddings = self.lm.get_input_embeddings()

    print("Restoring pretrained weights for the visual model.")
    if 'clip' in visual_encoder:
      self.visual_model = CLIPVisionModel.from_pretrained(visual_encoder)
    else:
      self.visual_model = AutoModel.from_pretrained(visual_encoder)

    if 'clip' in visual_encoder:
      hidden_size = self.visual_model.config.hidden_size
    else:
      raise NotImplementedError

    if self.args.freeze_vm:
      print("Freezing the VM.")
      self.visual_model.eval()
      for param in self.visual_model.parameters():
        param.requires_grad = False
    else:
      self.visual_model.train()

    self.visual_model_name = visual_encoder

    embedding_dim = self.input_embeddings.embedding_dim * self.args.n_visual_tokens
    self.text_hidden_fcs = nn.ModuleList([])
    if self.args.shared_emb_dim is None:
      if len(self.args.text_emb_layers) == 1:
        if (self.args.text_emb_layers[0] in [-1, self.lm.config.num_hidden_layers]) and ('bert' not in opt_version):
          out_dim = self.lm.config.word_embed_proj_dim
        else:
          out_dim = self.lm.config.hidden_size
      else:
        if (-1 in self.args.text_emb_layers) or (self.lm.config.num_hidden_layers in self.args.text_emb_layers) \
          and (self.lm.config.word_embed_proj_dim != self.lm.config.hidden_size):
          raise ValueError('No projection dim specified but model uses last output layer and an intermediate one (which have different dims).')
        else:
          out_dim = self.lm.config.hidden_size
    else:
      out_dim = self.args.shared_emb_dim

      for layer_idx in self.args.text_emb_layers:
        if (layer_idx == -1 or layer_idx == self.lm.config.num_hidden_layers) and ('bert' not in opt_version):
          in_dim = self.lm.config.word_embed_proj_dim

          text_fc = [nn.Linear(in_dim, out_dim), nn.Dropout(self.args.text_embed_dropout_prob)]
          self.text_hidden_fcs.append(nn.Sequential(*text_fc))

        elif layer_idx < self.lm.config.num_hidden_layers:
          text_fc = [nn.Linear(self.lm.config.hidden_size, out_dim), nn.Dropout(self.args.text_embed_dropout_prob)]
          self.text_hidden_fcs.append(nn.Sequential(*text_fc))
        else:
          raise ValueError(f'Embedding of layer {layer_idx} was requested but model only has {self.lm.config.num_hidden_layers} layers.')

    self.visual_embeddings = nn.Linear(hidden_size, embedding_dim)
    self.visual_fc = nn.Linear(hidden_size, out_dim)

    self.image_dropout = nn.Dropout(self.args.image_embed_dropout_prob)


  def get_visual_embs(self, pixel_values: torch.FloatTensor, mode: str = 'captioning'):
    if mode not in ['captioning', 'retrieval']:
      raise ValueError(f'mode should be one of ["caption", "retrieval"], got {mode} instead.')

    # Extract visual embeddings from the vision encoder.
    if 'clip' in self.visual_model_name:
      outputs = self.visual_model(pixel_values)
      encoder_outputs = outputs.pooler_output
    else:
      raise NotImplementedError

    # Use the correct fc based on function argument.
    if mode == 'captioning':
      visual_embs = self.visual_embeddings(encoder_outputs)  # (2, D * n_visual_tokens)
      visual_embs = torch.reshape(visual_embs, (visual_embs.shape[0], self.args.n_visual_tokens, -1))
    elif mode == 'retrieval':
      visual_embs = self.visual_fc(encoder_outputs)  # (2, D * n_visual_tokens)
      visual_embs = torch.reshape(visual_embs, (visual_embs.shape[0], 1, -1))
    else:
      raise NotImplementedError

    visual_embs = self.image_dropout(visual_embs)
    return visual_embs


  def train(self, mode=True):
    super(FromageModel, self).train(mode=mode)
    # Overwrite train() to ensure Frozen models remain frozen.
    if self.args.freeze_lm:
      self.lm.eval()
    if self.args.freeze_vm:
      self.visual_model.eval()


  def forward(
    self,
    pixel_values: torch.FloatTensor,
    labels: torch.LongTensor,
    caption_len: torch.LongTensor,
    mode: str = 'captioning',
    concat_captions: bool = False,
    input_prefix: Optional[str] = None,
    inference: bool = False,
  ):
    visual_embs = self.get_visual_embs(pixel_values, mode)

    batch_size, vis_seq_len, _ = visual_embs.shape  # vis_seq_len = n_visual_tokens
    if labels is not None:
      assert labels.shape[0] == batch_size, (visual_embs.shape, labels.shape)

    input_embs = self.input_embeddings(labels)  # (N, T, D)

    last_embedding_idx = caption_len - 1  # -1 to retrieve the token before the eos token

    if input_prefix is not None:
      prompt_ids = self.tokenizer(input_prefix, add_special_tokens=False, return_tensors="pt").input_ids
      prompt_ids = prompt_ids.to(visual_embs.device)
      prompt_embs = self.input_embeddings(prompt_ids)
      prompt_embs = prompt_embs.repeat(batch_size, 1, 1)
      assert prompt_embs.shape[0] == batch_size, prompt_embs.shape
      assert prompt_embs.shape[2] == input_embs.shape[2], prompt_embs.shape
      assert len(prompt_embs.shape) == 3, prompt_embs.shape

    if mode == 'captioning':
      # Concat to text embeddings.
      condition_seq_len = 0
      if input_prefix is None:
        # Just add visual embeddings.
        input_embs = torch.cat([visual_embs, input_embs], axis=1)
        last_embedding_idx += vis_seq_len
        condition_seq_len += vis_seq_len
        full_labels = torch.zeros(visual_embs.shape[:2], dtype=torch.int64).to(visual_embs.device) - 100
      else:
        # Add visual and prompt embeddings.
        prefix_embs = torch.cat([visual_embs, prompt_embs], axis=1)
        input_embs = torch.cat([prefix_embs, input_embs], axis=1)

        last_embedding_idx += prefix_embs.shape[1]
        condition_seq_len += prefix_embs.shape[1]
        full_labels = torch.zeros(prefix_embs.shape[:2], dtype=torch.int64).to(visual_embs.device) - 100

      # Mask out embedding tokens in the labels.
      full_labels = torch.cat([full_labels, labels], axis=1)

      pad_idx = []

      for label in full_labels:
        for k, token in enumerate(label):
          # Mask out retrieval token if it exists.
          if token in [self.tokenizer.pad_token_id, self.retrieval_token_idx]:
            label[k:] = -100
            pad_idx.append(k)
            break
          if k == len(label) - 1:  # No padding found.
            pad_idx.append(k + 1)
      assert len(pad_idx) == batch_size, (len(pad_idx), batch_size)

      bs, seq_len, embs_dim = input_embs.shape
      if concat_captions:
        assert len(input_embs.shape) == 3, input_embs
        assert len(full_labels.shape) == 2, full_labels
        assert batch_size % 2 == 0
        all_concat_input_embs = []
        all_concat_labels = []

        # Rearrange embeddings and labels (and their padding) to concatenate captions.
        for i in range(batch_size // 2):
          first_idx = i * 2
          second_idx = first_idx + 1
          first_emb = input_embs[first_idx, :pad_idx[first_idx], :]
          first_labels = full_labels[first_idx, :pad_idx[first_idx]]
          first_padding = input_embs[first_idx, pad_idx[first_idx]:, :]
          first_labels_padding = full_labels[first_idx, pad_idx[first_idx]:]

          second_emb = input_embs[second_idx, :pad_idx[second_idx], :]
          second_labels = full_labels[second_idx, :pad_idx[second_idx]]
          second_padding = input_embs[second_idx, pad_idx[second_idx]:, :]
          second_labels_padding = full_labels[second_idx, pad_idx[second_idx]:]

          assert torch.all(first_labels_padding == -100), first_labels_padding
          assert torch.all(second_labels_padding == -100), second_labels_padding
          concat_input_embs = torch.cat([first_emb, second_emb, first_padding, second_padding], axis=0)   # (T*2, 768)
          concat_labels = torch.cat([first_labels, second_labels, first_labels_padding, second_labels_padding], axis=0)   # (T*2, 768)
          all_concat_input_embs.append(concat_input_embs)
          all_concat_labels.append(concat_labels)

        # Pad to max length.
        input_embs = torch.stack(all_concat_input_embs, axis=0)  # (N/2, T*2, 768)
        full_labels = torch.stack(all_concat_labels, axis=0)  # (N/2, T*2, 768)
        assert input_embs.shape == (bs // 2, seq_len * 2, embs_dim), input_embs.shape
        assert full_labels.shape == (bs // 2, seq_len * 2), full_labels.shape

      output = self.lm(inputs_embeds=input_embs,
                       labels=full_labels,
                       output_hidden_states=True)
    elif mode == 'retrieval':
      full_labels = torch.clone(labels)
      if input_prefix is not None:
        print(f'Adding prefix "{input_prefix}" to retrieval.')
        # Add prompt embeddings.
        prefix_embs = prompt_embs
        input_embs = torch.cat([prefix_embs, input_embs], axis=1)
        last_embedding_idx += prefix_embs.shape[1]
        full_labels = torch.cat([
          torch.zeros(prefix_embs.shape[:2], dtype=torch.int64).to(labels.device) - 100,
          full_labels
        ], axis=1)
      
      pad_idx = []
      for label in full_labels:
        for k, token in enumerate(label):
          if token == self.tokenizer.pad_token_id:
            label[k:] = -100
            pad_idx.append(k)
            break
          if k == len(label) - 1:  # No padding found.
            pad_idx.append(k + 1)
      assert len(pad_idx) == batch_size, (len(pad_idx), batch_size)

      output = self.lm(inputs_embeds=input_embs,
                       labels=full_labels,
                       output_hidden_states=True)
    else:
      raise NotImplementedError

    last_embedding = None
    last_output_logit = None
    hidden_states = []

    if mode == 'retrieval':
      if self.args.shared_emb_dim is not None:
        for idx, fc_layer in zip(self.args.text_emb_layers, self.text_hidden_fcs):
          hidden_states.append(fc_layer(output.hidden_states[idx]))  # (N, seq_len, 2048)
      else:
        for idx in self.args.text_emb_layers:
          hidden_states.append(output.hidden_states[idx])

      # Add hidden states together.
      last_hidden_state = torch.stack(hidden_states, dim=-1).sum(dim=-1)

      if not concat_captions:
        last_embedding = torch.stack([last_hidden_state[i, last_embedding_idx[i], :] for i in range(batch_size)], axis=0)  # (N, D)
        last_output_logit = torch.stack([output.logits[i, last_embedding_idx[i] - 1, :] for i in range(batch_size)], axis=0)  # (N, D)
      else:
        # Concatenate two captioning examples together.
        all_last_embedding = []
        all_last_output_logit = []
        for i in range(batch_size // 2):
          first_last_embedding_idx, second_last_embedding_idx = all_last_embedding_idx[i]
          first_last_embedding = last_hidden_state[i, first_last_embedding_idx, :]  # (N, D)
          first_last_output_logit = output.logits[i, first_last_embedding_idx - 1, :]  # (N, D)
          second_last_embedding = last_hidden_state[i, second_last_embedding_idx, :]  # (N, D)
          second_last_output_logit = output.logits[i, second_last_embedding_idx - 1, :]  # (N, D)
          all_last_embedding.append(first_last_embedding)
          all_last_embedding.append(second_last_embedding)
          all_last_output_logit.append(first_last_output_logit)
          all_last_output_logit.append(second_last_output_logit)

        last_embedding = torch.stack(all_last_embedding)
        last_output_logit = torch.stack(all_last_output_logit)

      # Compute retrieval loss.
      assert visual_embs.shape[1] == 1, visual_embs.shape
      visual_embs = visual_embs[:, 0, :]
      visual_embs = visual_embs / visual_embs.norm(dim=1, keepdim=True)
      last_embedding = last_embedding / last_embedding.norm(dim=1, keepdim=True)

      # cosine similarity as logits
      logit_scale = self.logit_scale.exp()
      visual_embs = logit_scale * visual_embs
    elif mode == 'captioning':
      pass
    else:
      raise NotImplementedError

    return output, full_labels, last_embedding, last_output_logit, visual_embs

  def generate(self, embeddings = torch.FloatTensor, max_len: int = 32,
               temperature: float = 0.0, top_p: float = 1.0, min_word_tokens: int = 0,
               ret_scale_factor: float = 1.0, filter_value: float = -float('Inf')):
    """Runs greedy decoding and returns generated captions.

    Args:
      embeddings: Input condition that the model uses for autoregressive generation.
      max_len: Maximum number of tokens to generate.
      temperature: Used to modulate logit distribution.
      top_p: If set to < 1, the smallest set of tokens with highest probabilities that add up to top_p or higher are kept for generation.
      min_word_tokens: Minimum number of words to generate before allowing a [RET] output.
      ret_scale_factor: Proportion to scale [RET] token logits by. A higher value may increase the probability of the model generating [RET] outputs.
      filter_value: Value to assign to tokens that should never be generated.
    Outputs:
      out: (N, T) int32 sequence of output tokens.
      output_embeddings: (N, T, 256) sequence of text output embeddings.
    """
    self.lm.eval()

    with torch.no_grad():  # no tracking history
      batch_size, s, _ = embeddings.shape
      # init output with image tokens
      out = None
      past_key_values = None
      output_embeddings = []
      output_logits = []

      for i in range(max_len):
        if 'opt' in self.opt_version:
          output = self.lm(inputs_embeds=embeddings, use_cache=False, output_hidden_states=True)
        else:
          if i == 0:
            output = self.lm(inputs_embeds=embeddings, use_cache=True, past_key_values=None, output_hidden_states=True)
          else:
            output = self.lm(input_ids=out[:, -1:], use_cache=True, past_key_values=past_key_values, output_hidden_states=True)

        # Collect and sum the hidden states.
        hidden_states = []
        if self.args.shared_emb_dim is not None:
          for idx, fc_layer in zip(self.args.text_emb_layers, self.text_hidden_fcs):
            hidden_states.append(fc_layer(output.hidden_states[idx]))  # (N, seq_len, 2048)
        else:
          for idx in self.args.text_emb_layers:
            hidden_states.append(output.hidden_states[idx])
        # Add hidden states together.
        last_hidden_state = torch.stack(hidden_states, dim=-1).sum(dim=-1)  # (N, T, 256)
        last_embedding = last_hidden_state / last_hidden_state.norm(dim=-1, keepdim=True)
        output_embeddings.append(last_embedding)

        logits = output.logits[:, -1, :]  # (N, vocab_size)
        if top_p == 1.0:
          logits = logits.cpu()
        output_logits.append(logits)

        if self.retrieval_token_idx != -1 and self.retrieval_token_idx is not None:
          if i < min_word_tokens:
            # Eliminate probability of generating [RET] if this is earlier than min_word_tokens.
            logits[:, self.retrieval_token_idx] = filter_value
          else:
            # Multiply by scaling factor.
            logits[:, self.retrieval_token_idx] = logits[:, self.retrieval_token_idx] * ret_scale_factor

        past_key_values = output.past_key_values

        if temperature == 0.0:
          if top_p != 1.0:
            raise ValueError('top_p cannot be set if temperature is 0 (greedy decoding).')
          next_token = torch.argmax(logits, keepdim=True, dim=-1)  # (N, 1)
        else:
          logits = logits / temperature

          # Apply top-p filtering.
          if top_p < 1.0:
            assert top_p > 0, f'top_p should be above 0, got {top_p} instead.'
            sorted_logits, sorted_indices = torch.sort(logits, descending=True)  # (N, D) and (N, D)
            cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) # (N, D)

            # Remove tokens with cumulative probability above the threshold
            sorted_indices_to_remove = cumulative_probs > top_p
            # Shift the indices to the right to keep also the first token above the threshold
            sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
            sorted_indices_to_remove[..., 0] = 0

            for j in range(sorted_indices.shape[0]):
              indices_to_remove = sorted_indices[j, sorted_indices_to_remove[j, :]]
              logits[j, indices_to_remove] = filter_value

          token_weights = logits.exp()   # (N, vocab_size)
          next_token = torch.multinomial(token_weights, 1)  # (N, 1)

        next_token = next_token.long().to(embeddings.device)
        if out is not None:
          out = torch.cat([out, next_token], dim=-1)
        else:
          out = next_token

        if 'opt' in self.opt_version:
          next_embedding = self.input_embeddings(next_token)
          embeddings = torch.cat([embeddings, next_embedding], dim=1)
        elif (self.tokenizer.eos_token_id and (next_token == self.tokenizer.eos_token_id).all()):
          # End of generation.
          break

    return out, output_embeddings, output_logits


class Fromage(nn.Module):
  def __init__(self, tokenizer, model_args: Optional[FrozenArgs] = None,
               path_array: Optional[List[str]] = None, emb_matrix: Optional[torch.tensor] = None):
    super().__init__()
    self.model = FromageModel(tokenizer, model_args)
    self.path_array = path_array
    self.emb_matrix = emb_matrix

  def __call__(self, images: Tensor, tgt_tokens: Optional[Tensor] = None, caption_len: Optional[Tensor] = None,
               generate: bool = False, num_words: int = 32, temperature: float = 1.0, top_p: float = 1.0,
               ret_scale_factor: float = 1.0, min_word_tokens: int = 0,
               mode: str = 'captioning', concat_captions: bool = False,
               input_prefix: Optional[str] = None, inference: bool = False) -> Tensor:
    if generate:
      return self.model.generate(images, num_words, temperature=temperature, top_p=top_p,
                                 min_word_tokens=min_word_tokens, ret_scale_factor=ret_scale_factor)
    else:
      output = self.model(
        pixel_values = images,
        labels = tgt_tokens,
        caption_len = caption_len,
        mode = mode,
        concat_captions = concat_captions,
        input_prefix = input_prefix,
        inference = inference)
      return output

  def generate_for_images_and_texts(
    self, prompts: List, num_words: int = 0, ret_scale_factor: float = 1.0, top_p: float = 1.0, temperature: float = 0.0,
    max_num_rets: int = 1):
    """
    Encode prompts into embeddings.

    Args:
      prompts: List of interleaved PIL.Image.Image and strings representing input to the model.
      num_words: Maximum number of words to generate for. If num_words = 0, the model will run its forward pass and return the outputs.
      ret_scale_factor: Proportion to scale [RET] token logits by. A higher value may increase the probability of the model generating [RET] outputs.
      top_p: If set to < 1, the smallest set of tokens with highest probabilities that add up to top_p or higher are kept for generation.
      temperature: Used to modulate logit distribution.
      max_num_rets: Maximum number of images to return in one generation pass.
    Returns:
      return_outputs: List consisting of either str or List[PIL.Image.Image] objects, representing image-text interleaved model outputs.
    """
    input_embs = []
    input_ids = []
    add_bos = True

    for i, p in enumerate(prompts):
      if type(p) == Image.Image:
        # Encode as image.
        pixel_values = utils.get_pixel_values_for_model(self.model.feature_extractor, p)
        pixel_values = pixel_values.to(device=self.model.logit_scale.device, dtype=self.model.logit_scale.dtype)
        pixel_values = pixel_values[None, ...]

        visual_embs = self.model.get_visual_embs(pixel_values, mode='captioning')  # (1, n_visual_tokens, D)
        input_embs.append(visual_embs)
      elif type(p) == str:
        text_ids = self.model.tokenizer(p, add_special_tokens=True, return_tensors="pt").input_ids.to(self.model.logit_scale.device)
        if not add_bos:
          # Remove <bos> tag.
          text_ids = text_ids[:, 1:]
        else:
          # Only add <bos> once.
          add_bos = False

        text_embs = self.model.input_embeddings(text_ids)  # (1, T, D)
        input_embs.append(text_embs)
        input_ids.append(text_ids)
      else:
        raise ValueError(f'Input prompts should be either PIL.Image.Image or str types, got {type(p)} instead.')
    input_embs = torch.cat(input_embs, dim=1)
    input_ids = torch.cat(input_ids, dim=1)

    if num_words == 0:
      generated_ids = input_ids
      outputs = self.model.lm(inputs_embeds=input_embs, use_cache=False, output_hidden_states=True)
      # Map outputs to embeddings, so we can retrieve embeddings from the [RET] tokens.
      out = []
      for x, fc in zip(self.model.args.text_emb_layers, self.model.text_hidden_fcs):
          out.append(fc(outputs.hidden_states[x]))
      embeddings = torch.stack(out, dim=-1).sum(dim=-1)
      embeddings = embeddings / embeddings.norm(dim=-1, keepdim=True)  # (N, T, 256)
    elif num_words > 0:
      generated_ids, generated_embeddings, _ = self.model.generate(input_embs, num_words,
        temperature=temperature, top_p=top_p, ret_scale_factor=ret_scale_factor)
      embeddings = generated_embeddings[-1][:, input_embs.shape[1]:]

      # Truncate to newline.
      newline_token_id = self.model.tokenizer('\n', add_special_tokens=False).input_ids[0]
      trunc_idx = 0
      for j in range(generated_ids.shape[1]):
        if generated_ids[0, j] == newline_token_id:
          trunc_idx = j
          break
      if trunc_idx > 0:
        generated_ids = generated_ids[:, :trunc_idx]
        embeddings = embeddings[:, :trunc_idx]
    else:
      raise ValueError

    # Save outputs as an interleaved list.
    return_outputs = []
    # Find up to max_num_rets [RET] tokens, and their corresponding scores.
    all_ret_idx = [i for i, x in enumerate(generated_ids[0, :] == self.model.retrieval_token_idx) if x][:max_num_rets]
    seen_image_idx = []  # Avoid showing the same image multiple times.

    last_ret_idx = 0
    if len(all_ret_idx) == 0:
      # No [RET] tokens.
      caption = self.model.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
      return_outputs.append(utils.truncate_caption(caption))
    else:
      for ret_idx in all_ret_idx:
        ret_emb = embeddings[:, ret_idx, :]
        scores = self.emb_matrix @ ret_emb.T

        # Downweight seen images.
        for seen_idx in seen_image_idx:
          scores[seen_idx, :] -= 1000

        # Get the top 3 images for each image.
        _, top_image_idx = scores.squeeze().topk(3)
        image_outputs = []
        for img_idx in top_image_idx:
          # Find the first image that does not error out.
          try:
            seen_image_idx.append(img_idx)
            img = utils.get_image_from_url(self.path_array[img_idx])
            image_outputs.append(img)
            if len(image_outputs) == max_num_rets:
              break
          except UnidentifiedImageError:
            pass

        caption = self.model.tokenizer.batch_decode(generated_ids[:, last_ret_idx:ret_idx], skip_special_tokens=True)[0]
        last_ret_idx = ret_idx + 1
        return_outputs.append(utils.truncate_caption(caption) + ' [RET]')
        return_outputs.append(image_outputs)

    return return_outputs


def load_fromage(model_dir: str, ckpt_path: str) -> Fromage:
  model_args_path = os.path.join(model_dir, 'model_args.json')
  model_ckpt_path = os.path.join(ckpt_path)
  embs_paths = [s for s in glob.glob(os.path.join(model_dir, 'cc3m_embeddings*.pkl'))]

  if not os.path.exists(model_args_path):
    raise ValueError(f'model_args.json does not exist in {model_dir}.')
  if not os.path.exists(model_ckpt_path):
    raise ValueError(f'pretrained_ckpt.pth.tar does not exist in {model_dir}.')
  if len(embs_paths) == 0:
    raise ValueError(f'cc3m_embeddings_*.pkl files do not exist in {model_dir}.')

  # Load embeddings.
  # Construct embedding matrix for nearest neighbor lookup.
  path_array = []
  emb_matrix = []

  # These were precomputed for all CC3M images with `model.get_visual_embs(image, mode='retrieval')`.
  for p in embs_paths:
    with open(p, 'rb') as wf:
        train_embs_data = pkl.load(wf)
        path_array.extend(train_embs_data['paths'])
        emb_matrix.append(train_embs_data['embeddings'])
  emb_matrix = np.concatenate(emb_matrix, axis=0)

  # Number of paths should be equal to number of embeddings.
  assert len(path_array) == emb_matrix.shape[0], (len(path_array), emb_matrix.shape[0])

  with open(model_args_path, 'r') as f:
      model_kwargs = json.load(f)

  # Initialize tokenizer.
  tokenizer = GPT2Tokenizer.from_pretrained(model_kwargs['opt_version'])
  tokenizer.pad_token = tokenizer.eos_token
  # Add special tokens to the model to enable [RET].
  tokenizer.add_special_tokens({"cls_token": "<|image|>"})
  tokenizer.add_tokens('[RET]')
  ret_token_idx = tokenizer('[RET]', add_special_tokens=False).input_ids
  assert len(ret_token_idx) == 1, ret_token_idx
  model_kwargs['retrieval_token_idx'] = ret_token_idx[0]
  args = namedtuple('args', model_kwargs)(**model_kwargs)

  # Initialize model for inference.
  model = Fromage(tokenizer, args, path_array=path_array, emb_matrix=emb_matrix)
  model = model.eval()
  model = model.bfloat16()
  model = model.cuda()

  # Load pretrained linear mappings and [RET] embeddings.
  checkpoint = torch.load(model_ckpt_path)
  model.load_state_dict(checkpoint['state_dict'], strict=False)
  with torch.no_grad():
      model.model.input_embeddings.weight[model.model.retrieval_token_idx, :].copy_(checkpoint['state_dict']['ret_input_embeddings.weight'].cpu().detach())

  logit_scale = model.model.logit_scale.exp()
  emb_matrix = torch.tensor(emb_matrix, dtype=logit_scale.dtype).to(logit_scale.device)
  emb_matrix = emb_matrix / emb_matrix.norm(dim=1, keepdim=True)
  emb_matrix = logit_scale * emb_matrix
  model.emb_matrix = emb_matrix

  return model