1. 1
2. 2
3. 3
4. 4
5. 5
6. 6
7. 7
8. 8
9. 9
10. 10
11. 11
12. 12
13. 13
14. 14
15. 15
16. 16
17. 17
18. 18
19. 19
20. 20
21. 21
22. 22
23. 23
24. 24
25. 25
26. 26
27. 27
28. 28
29. 29
30. 30
31. 31
32. 32
33. 33
34. 34
35. 35
36. 36
37. 37
38. 38

import torch
from imagebind import data
from imagebind.models import imagebind_model
from imagebind.models.imagebind_model import ModalityType

device = "cuda:0" if torch.cuda.is_available() else "cpu"

print("Loading model")
model = imagebind_model.imagebind_huge(pretrained=True)
model.eval()
model.to(device)


def embed(mode, data):
    with torch.no_grad():
        return model({mode: data})[mode][0]


def embed_text(text):
    return embed(ModalityType.TEXT, data.load_and_transform_text([text], device))


def embed_audio(audio_path):
    return embed(
        ModalityType.AUDIO, data.load_and_transform_audio_data([audio_path], device)
    )


def embed_image(image_path):
    return embed(
        ModalityType.VISION, data.load_and_transform_vision_data([image_path], device)
    )


def embed_video(video_path):
    return embed(
        ModalityType.VISION, data.load_and_transform_video_data([video_path], device)
    )