Class OpenClipEncoder

extra/models/clip.py:428–480

}

class OpenClipEncoder:
  def __init__(self, dims:int, text_cfg:Dict, vision_cfg:Dict, **_):
    self.visual = Open.ClipVisionTransformer(**vision_cfg)

    text = Open.ClipTextTransformer(**text_cfg)
    self.transformer = text.transformer
    self.token_embedding = text.token_embedding
    self.positional_embedding = text.positional_embedding
    self.ln_final = text.ln_final
    self.text_projection = text.text_projection

    self.attn_mask = Tensor.full((77, 77), float("-inf")).triu(1).realize()
    self.mean = Tensor([0.48145466, 0.45782750, 0.40821073]).reshape(-1, 1, 1)
    self.std  = Tensor([0.26862954, 0.26130258, 0.27577711]).reshape(-1, 1, 1)

  # TODO:
  # Should be doable in pure tinygrad, would just require some work and verification.
  # This is very desirable since it would allow for full generation->evaluation in a single JIT call.
  def prepare_image(self, image) -> Tensor:
    from PIL import Image
    SIZE = 224
    w, h = image.size
    scale = min(SIZE / h, SIZE / w)
    image = image.resize((max(int(w*scale),SIZE),max(int(h*scale),SIZE)), Image.Resampling.BICUBIC)
    w, h = image.size
    if w > SIZE:
      left = (w - SIZE) // 2
      image = image.crop((left, left+SIZE, 0, SIZE))
    elif h > SIZE:
      top = (h - SIZE) // 2
      image = image.crop((0, SIZE, top, top+SIZE))

    x = Tensor(np.array(image.convert('RGB')), device=self.std.device)
    x = x.permute(2, 0, 1).cast(dtypes.float32) / 255.0
    return (x - self.mean) / self.std

  def encode_tokens(self, tokens:Tensor) -> Tensor:
    x = self.token_embedding(tokens)
    x = x + self.positional_embedding
    x = self.transformer(x, attn_mask=self.attn_mask)
    x = self.ln_final(x)
    x = x[Tensor.arange(x.shape[0], device=x.device), tokens.argmax(axis=-1)]
    x = x @ self.text_projection
    return x

  def get_clip_score(self, tokens:Tensor, image:Tensor) -> Tensor:
    image_features: Tensor = self.visual(image)
    image_features /= image_features.square().sum(-1, keepdim=True).sqrt() # Frobenius Norm

    text_features = self.encode_tokens(tokens)
    text_features /= text_features.square().sum(-1, keepdim=True).sqrt() # Frobenius Norm

    return (image_features * text_features).sum(axis=-1)