| import os |
| import torch |
| import numpy as np |
| from typing import BinaryIO, List |
|
|
| from imagebind import imagebind_model |
| from imagebind.models.imagebind_model import ModalityType |
| from imagebind.models.multimodal_preprocessors import SimpleTokenizer, TextPreprocessor |
|
|
|
|
| V2_URL = "https://huggingface.co/jondurbin/videobind-v0.2/resolve/main/videobind.pth" |
| V2_PATH = "./.checkpoints/videobind-v0.2.pth" |
| BPE_PATH = "./models/bpe_simple_vocab_16e6.txt.gz" |
| TOKENIZER = SimpleTokenizer(bpe_path=BPE_PATH) |
| LENGTH_TOKENIZER = SimpleTokenizer(bpe_path=BPE_PATH, context_length=1024) |
| TOKEN_CHUNK_SIZE = 74 |
|
|
| def get_imagebind_v2(path: str=V2_PATH): |
| if not os.path.isfile(path): |
| os.makedirs(os.path.dirname(path), exist_ok=True) |
| torch.hub.download_url_to_file(V2_URL, path, progress=True) |
| imagebind_model = torch.load(path) |
| return imagebind_model |
|
|
|
|
| def load_and_transform_text(text, device): |
| if text is None: |
| return None |
| tokens = [TOKENIZER(t).unsqueeze(0).to(device) for t in text] |
| tokens = torch.cat(tokens, dim=0) |
| return tokens |
|
|
| def split_text_by_token_limit(text, tokenizer, max_tokens=TOKEN_CHUNK_SIZE): |
| def fits_in_token_limit(text_segment): |
| tokens = tokenizer(text_segment) |
| tokens = tokens[tokens != 0][1:-1].tolist() |
| return len(tokens) <= max_tokens |
|
|
| def recursive_split(text, delimiters): |
| if fits_in_token_limit(text): |
| return [text] |
| if not delimiters: |
| return split_by_tokens(text) |
| delimiter = delimiters[0] |
| parts = text.split(delimiter) |
| result = [] |
| current_segment = "" |
| for part in parts: |
| candidate_segment = current_segment + (delimiter if current_segment else '') + part |
| if fits_in_token_limit(candidate_segment): |
| current_segment = candidate_segment |
| else: |
| if current_segment: |
| result.append(current_segment) |
| current_segment = part |
| if current_segment: |
| result.append(current_segment) |
| final_result = [] |
| for segment in result: |
| if fits_in_token_limit(segment): |
| final_result.append(segment) |
| else: |
| final_result.extend(recursive_split(segment, delimiters[1:])) |
| return final_result |
|
|
| def split_by_tokens(text): |
| tokens = tokenizer(text) |
| tokens = tokens[tokens != 0][1:-1].tolist() |
| chunks = np.array_split(tokens, int(len(tokens) / max_tokens) or 1) |
| return [ |
| tokenizer.decode(segment_tokens) |
| for segment_tokens in chunks |
| ] |
|
|
| return recursive_split(text, ['\n', '.', '!', '?', ',', ' ']) |
|
|
| def load_and_transform_text_chunks(text, device): |
| if not text: |
| return [] |
| all_tokens = LENGTH_TOKENIZER(text) |
| all_tokens = all_tokens[all_tokens != 0][1:-1].tolist() |
|
|
| return [ |
| load_and_transform_text([segment], device) |
| for segment in split_text_by_token_limit(text, LENGTH_TOKENIZER) |
| ] |
|
|
|
|
| class ImageBind: |
| def __init__(self, device="cuda:0", v2=False): |
| self.device = device |
| self.v2 = v2 |
| if v2: |
| if not os.path.exists(V2_PATH): |
| os.makedirs(os.path.dirname(V2_PATH), exist_ok=True) |
| torch.hub.download_url_to_file( |
| V2_URL, |
| V2_PATH, |
| progress=True, |
| ) |
| self.imagebind = torch.load(V2_PATH) |
| else: |
| self.imagebind = imagebind_model.imagebind_huge(pretrained=True) |
| self.imagebind.eval() |
| self.imagebind.to(self.device) |
|
|
| def generate_text_embeddings(self, text: str): |
| if not self.v2: |
| return self.imagebind({ |
| ModalityType.TEXT: load_and_transform_text([text], self.device) |
| })[ModalityType.TEXT] |
| chunks = load_and_transform_text_chunks(text, self.device) |
| embeddings = [ |
| self.imagebind({ModalityType.TEXT: chunk})[ModalityType.TEXT] |
| for chunk in chunks |
| ] |
| return torch.mean(torch.stack(embeddings), dim=0) |
|
|
| """ Deactivating full embeddings as they are not used in the current implementation |
| def get_inputs(self, video_file: BinaryIO) -> dict: |
| audio_file = video_utils.copy_audio(video_file.name) |
| try: |
| duration = video_utils.get_video_duration(video_file.name) |
| video_data = data.load_and_transform_video_data( |
| [video_file.name], |
| self.device, |
| ) |
| audio_data = data.load_and_transform_audio_data( |
| [audio_file.name], |
| self.device, |
| ) |
| inputs = { |
| ModalityType.VISION: video_data, |
| ModalityType.AUDIO: audio_data, |
| } |
| return inputs |
| finally: |
| audio_file.close() |
| |
| @torch.no_grad() |
| def embed(self, descriptions: List[str], video_files: List[BinaryIO]) -> Embeddings: |
| return_value = None |
| for idx in range(len(descriptions)): |
| inputs = self.get_inputs(video_files[idx]) |
| embeddings = self.imagebind(inputs) |
| text_embeddings = self.generate_text_embeddings(descriptions[idx]) |
| if not return_value: |
| return_value = Embeddings( |
| video=embeddings[ModalityType.VISION], |
| audio=embeddings[ModalityType.AUDIO], |
| description=text_embeddings, |
| ) |
| else: |
| return_value.video = torch.cat((return_value.video, embeddings[ModalityType.VISION])) |
| return_value.audio = torch.cat((return_value.audio, embeddings[ModalityType.AUDIO])) |
| return_value.description = torch.cat((return_value.description, text_embeddings)) |
| return return_value |
| |
| @torch.no_grad() |
| def embed_only_video(self, video_files: List[BinaryIO]) -> Embeddings: |
| video_filepaths = [video_file.name for video_file in video_files] |
| durations = [video_utils.get_video_duration(f.name) for f in video_files] |
| embeddings = self.imagebind({ |
| ModalityType.VISION: [ |
| data.load_and_transform_video_data( |
| [video_filepaths[idx]], |
| self.device, |
| )[0] |
| for idx in range(len(video_filepaths)) |
| ] |
| }) |
| return Embeddings( |
| video=embeddings[ModalityType.VISION], |
| ) |
| |
| @torch.no_grad() |
| def embed_video_and_text(self, video_files: List[BinaryIO], descriptions: List[str]) -> Embeddings: |
| video_filepaths = [video_file.name for video_file in video_files] |
| durations = [video_utils.get_video_duration(f.name) for f in video_files] |
| embeddings = self.imagebind({ |
| ModalityType.VISION: [ |
| data.load_and_transform_video_data( |
| [video_filepaths[idx]], |
| self.device, |
| )[0] |
| for idx in range(len(video_filepaths)) |
| ], |
| }) |
| description_embeddings = torch.stack([ |
| self.generate_text_embeddings(description) |
| for description in descriptions |
| ]) |
| return Embeddings( |
| video=embeddings[ModalityType.VISION], |
| description=description_embeddings, |
| ) |
| |
| @torch.no_grad() |
| def embed_text(self, texts: List[str]) -> torch.Tensor: |
| return_value = None |
| for text in texts: |
| emb = self.generate_text_embeddings(text) |
| if not return_value: |
| return_value = emb |
| else: |
| return_value = torch.cat((return_value, emb)) |
| return return_value |
| """ |
| |
| @torch.no_grad() |
| def embed_text(self, texts: List[str]) -> torch.Tensor: |
| embeddings = [] |
| for text in texts: |
| emb = self.generate_text_embeddings(text) |
| embeddings.append(emb) |
| |
| if not embeddings: |
| return None |
| |
| |
| return torch.stack(embeddings, dim=0) |
|
|
