[ENH]: OpenCLIP EF device param

[tazarov]

Ref: https://discord.com/channels/1073293645303795742/1214028592372252682

Description of changes

Summarize the changes made by this PR.

• Improvements & Bug fixes
  • Added device optional param to OpenCLIP EF

Test plan

How are these changes tested?

• Tests pass locally with pytest for python, yarn test for js

Documentation Changes

NOTE: It doesn't see we have OpenCLIP EF docs. Will have to add.

[github-actions]

Reviewer Checklist

Please leverage this checklist to ensure your code review is thorough before approving

Testing, Bugs, Errors, Logs, Documentation

• Can you think of any use case in which the code does not behave as intended? Have they been tested?
• Can you think of any inputs or external events that could break the code? Is user input validated and safe? Have they been tested?
• If appropriate, are there adequate property based tests?
• If appropriate, are there adequate unit tests?
• Should any logging, debugging, tracing information be added or removed?
• Are error messages user-friendly?
• Have all documentation changes needed been made?
• Have all non-obvious changes been commented?

System Compatibility

• Are there any potential impacts on other parts of the system or backward compatibility?
• Does this change intersect with any items on our roadmap, and if so, is there a plan for fitting them together?

Quality

• Is this code of a unexpectedly high quality (Readability, Modularity, Intuitiveness)

[brunorosilva]

Is this already working? It seems that when computing embeddings, the model is on GPU but the data is on CPU (code below).

from chromadb.utils.embedding_functions import OpenCLIPEmbeddingFunction
from chromadb.utils.data_loaders import ImageLoader
import os
import chromadb
from tqdm import tqdm
import numpy as np
from PIL import Image
embedding_function = OpenCLIPEmbeddingFunction(device='cuda')
data_loader = ImageLoader()

client = chromadb.PersistentClient(path="./.chroma")
collection = client.create_collection(
    name='multimodal_collection', 
    embedding_function=embedding_function)
CHUNK_SIZE = 1
for idx in tqdm(range(0, len(os.listdir("wikiart")[:2]), CHUNK_SIZE)):
    ids = [str(i) for i in range(len(os.listdir("wikiart")))[idx:idx+CHUNK_SIZE]]
    imgs = [np.array(Image.open(f"wikiart/{img_name}")) for img_name in os.listdir("wikiart")[idx:idx+CHUNK_SIZE]]
    collection.add(ids=ids, images=imgs)

Error:

Traceback (most recent call last):
  File "/home/rosilva/Desktop/projects/art-img/embed_wikiart.py", line 26, in <module>
    collection.add(ids=ids, images=imgs)
  File "/home/rosilva/.cache/pypoetry/virtualenvs/art-img-zVHhaY_I-py3.10/lib/python3.10/site-packages/chromadb/api/models/Collection.py", line 156, in add
    embeddings = self._embed(input=images)
  File "/home/rosilva/.cache/pypoetry/virtualenvs/art-img-zVHhaY_I-py3.10/lib/python3.10/site-packages/chromadb/api/models/Collection.py", line 633, in _embed
    return self._embedding_function(input=input)
  File "/home/rosilva/.cache/pypoetry/virtualenvs/art-img-zVHhaY_I-py3.10/lib/python3.10/site-packages/chromadb/api/types.py", line 193, in __call__
    result = call(self, input)
  File "/home/rosilva/.cache/pypoetry/virtualenvs/art-img-zVHhaY_I-py3.10/lib/python3.10/site-packages/chromadb/utils/embedding_functions.py", line 739, in __call__
    embeddings.append(self._encode_image(cast(Image, item)))
  File "/home/rosilva/.cache/pypoetry/virtualenvs/art-img-zVHhaY_I-py3.10/lib/python3.10/site-packages/chromadb/utils/embedding_functions.py", line 723, in _encode_image
    image_features = self._model.encode_image(
  File "/home/rosilva/.cache/pypoetry/virtualenvs/art-img-zVHhaY_I-py3.10/lib/python3.10/site-packages/open_clip/model.py", line 266, in encode_image
    features = self.visual(image)
  File "/home/rosilva/.cache/pypoetry/virtualenvs/art-img-zVHhaY_I-py3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/rosilva/.cache/pypoetry/virtualenvs/art-img-zVHhaY_I-py3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/rosilva/.cache/pypoetry/virtualenvs/art-img-zVHhaY_I-py3.10/lib/python3.10/site-packages/open_clip/transformer.py", line 503, in forward
    x = self.conv1(x)  # shape = [*, width, grid, grid]
  File "/home/rosilva/.cache/pypoetry/virtualenvs/art-img-zVHhaY_I-py3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/rosilva/.cache/pypoetry/virtualenvs/art-img-zVHhaY_I-py3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/rosilva/.cache/pypoetry/virtualenvs/art-img-zVHhaY_I-py3.10/lib/python3.10/site-packages/torch/nn/modules/conv.py", line 460, in forward
    return self._conv_forward(input, self.weight, self.bias)
  File "/home/rosilva/.cache/pypoetry/virtualenvs/art-img-zVHhaY_I-py3.10/lib/python3.10/site-packages/torch/nn/modules/conv.py", line 456, in _conv_forward
    return F.conv2d(input, weight, bias, self.stride,
RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same or input should be a MKLDNN tensor and weight is a dense tensor

This should be fixed by loading the batch to GPU when creating the embedding. I can work on this, but wanted to make sure I'm not doing anything silly beforehand.

[MartinFix]

Same issue here, specifically for the OpenClipEmbeddingModel. I ended up by "patching" the file open_clip\tokenizer.py for the class SimpleTokenizer in function __call__ like this and it worked:
result = torch.zeros(len(all_tokens), context_length, dtype=torch.long).to('cuda')

This is surely only a brute-force hack for one particular case and I'm pretty sure there would be a more elegant and generally applicable solution ;-)