Update app.py

app.py

@@ -4,6 +4,7 @@ import spaces
 import torch
 
 from PIL import Image
+from decord import VideoReader, cpu
 from transformers import AutoModel, AutoTokenizer
 
 # Pre-Initialize
@@ -14,6 +15,7 @@ print(f"[SYSTEM] | Using {DEVICE} type compute device.")
 
 # Variables
 DEFAULT_INPUT = "Describe in one paragraph."
+MAX_FRAMES = 64
 
 repo = AutoModel.from_pretrained("openbmb/MiniCPM-V-2_6", torch_dtype=torch.bfloat16, trust_remote_code=True)
 tokenizer = AutoTokenizer.from_pretrained("openbmb/MiniCPM-V-2_6", trust_remote_code=True)
@@ -27,6 +29,21 @@ footer {
 '''
 
 # Functions
+def encode_video(video_path):
+    def uniform_sample(l, n):
+        gap = len(l) / n
+        idxs = [int(i * gap + gap / 2) for i in range(n)]
+        return [l[i] for i in idxs]
+
+    vr = VideoReader(video_path, ctx=cpu(0))
+    sample_fps = round(vr.get_avg_fps() / 1)
+    frame_idx = [i for i in range(0, len(vr), sample_fps)]
+    if len(frame_idx) > MAX_NUM_FRAMES:
+        frame_idx = uniform_sample(frame_idx, MAX_FRAMES)
+    frames = vr.get_batch(frame_idx).asnumpy()
+    frames = [Image.fromarray(v.astype('uint8')) for v in frames]
+    return frames
+    
 @spaces.GPU(duration=60)
 def generate(image, video, instruction=DEFAULT_INPUT, sampling=False, temperature=0.7, top_p=0.8, top_k=100, repetition_penalty=1.05, max_tokens=512):
     repo.to(DEVICE)
@@ -36,11 +53,10 @@ def generate(image, video, instruction=DEFAULT_INPUT, sampling=False, temperatur
     
     if not video:
         image_data = Image.fromarray(image.astype('uint8'), 'RGB')
-        print(image_data, instruction)
-
         inputs = [{"role": "user", "content": [image_data, instruction]}]
     else:
-        priny("video")
+        video_data = encode_video(video)
+        inputs = [{"role": "user", "content": video_data + [instruction]}]
     
     parameters = {
         "sampling": sampling,
@@ -49,6 +65,8 @@ def generate(image, video, instruction=DEFAULT_INPUT, sampling=False, temperatur
         "top_k": top_k,
         "repetition_penalty": repetition_penalty,
         "max_new_tokens": max_tokens
+        "use_image_id": False,
+        "max_slice_nums": 2,
     }
     
     output = repo.chat(image=None, msgs=inputs, tokenizer=tokenizer, **parameters)