File size: 1,185 Bytes
58ac08a 2e8783e 58ac08a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 |
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)
|