| import torch | |
| import numpy as np | |
| from transformers import AutoProcessor, AutoModel | |
| import decord | |
| class InternVLEmbedder: | |
| def __init__(self): | |
| self.device = "cuda" if torch.cuda.is_available() else "cpu" | |
| self.model = AutoModel.from_pretrained("OpenGVLab/InternVL2_5-1B-MPO", trust_remote_code=True).to(self.device) | |
| self.processor = AutoProcessor.from_pretrained("OpenGVLab/InternVL2_5-1B-MPO", trust_remote_code=True) | |
| def embed_video(self, video_path): | |
| vr = decord.VideoReader(video_path) | |
| frames = np.stack([vr[i].asnumpy() for i in np.linspace(0, len(vr)-1, 8).astype(int)]) | |
| tensor = torch.tensor(frames).permute(0, 3, 1, 2).unsqueeze(0).to(self.device) | |
| with torch.no_grad(): | |
| video_vector = self.model.get_video_features(tensor).squeeze(0).cpu().numpy() | |
| return video_vector / np.linalg.norm(video_vector) | |
| def embed_text(self, text): | |
| inputs = self.processor(text=[text], return_tensors="pt").to(self.device) | |
| with torch.no_grad(): | |
| text_vector = self.model.get_text_features(**inputs).squeeze(0).cpu().numpy() | |
| return text_vector / np.linalg.norm(text_vector) | |