Update app.py

app.py

@@ -32,10 +32,10 @@ footer {
 '''
 
 input_prefixes = {
-    "Image": "(A image file called █ has been attached) ",
-    "GIF": "(A GIF file called █ has been attached) ",
-    "Video": "(A video with audio file called █ has been attached) ",
-    "Audio": "(A audio file called █ has been attached) ",
+    "Image": "(A image file called █ has been attached, describe the image content) ",
+    "GIF": "(A GIF file called █ has been attached, describe the GIF content) ",
+    "Video": "(A video with audio file called █ has been attached, describe the video content and the audio content embedded into the video) ",
+    "Audio": "(A audio file called █ has been attached, describe the audio content) ",
 }
 
 filetypes = {
@@ -49,6 +49,18 @@ def uniform_sample(idxs, n):
     gap = len(idxs) / n
     return [idxs[int(i * gap + gap / 2)] for i in range(n)]
 
+def build_omni_chunks(path, sr=16000, seconds_per_unit=1):
+    vr = VideoReader(path, ctx=cpu(0))
+    fps = round(vr.get_avg_fps())
+    audio_np, _ = librosa.load(path, sr=sr, mono=True)
+    total_units = math.ceil(len(vr) / fps / seconds_per_unit)
+    content = []
+    for i in range(total_units):
+        frame = Image.fromarray(vr[int(i * fps * seconds_per_unit)].asnumpy().astype("uint8"))
+        audio_chunk = audio_np[sr * i * seconds_per_unit : sr * (i + 1) * seconds_per_unit]
+        content.extend(["<unit>", frame, audio_chunk])
+    return content
+    
 def encode_video(path):
     vr = VideoReader(path, ctx=cpu(0))
     fps = round(vr.get_avg_fps())
@@ -88,21 +100,11 @@ def generate(input, instruction=DEFAULT_INPUT, sampling=False, temperature=0.7,
         frames = encode_gif(input)
         content.extend(frames)
     elif filetype == "Video":
-        vr = VideoReader(input, ctx=cpu(0))
-        fps = round(vr.get_avg_fps())
-        raw_idxs = list(range(0, len(vr), fps))
-        idxs = raw_idxs if len(raw_idxs) <= MAX_FRAMES else uniform_sample(raw_idxs, MAX_FRAMES)
-        frames_np = vr.get_batch(idxs).asnumpy()
-        audio_np, sample_rate = librosa.load(input, sr=16000, mono=True)
-        for idx, frame_np in zip(idxs, frames_np):
-            image = Image.fromarray(frame_np.astype("uint8")).convert("RGB")
-            content.append(image)
-            sec = idx // fps
-            start = sec * sample_rate
-            end = start + sample_rate
-            chunk_np = audio_np[start:end]
-            chunk_tensor = torch.from_numpy(chunk_np).float().to(DEVICE)
-            content.append({"array": chunk_tensor, "sampling_rate": sample_rate})
+        omni_content = build_omni_chunks(input) + [instruction]
+        sys_msg = repo.get_sys_prompt(mode="omni", language="en")
+        msgs = [sys_msg, {"role": "user", "content": omni_content}]
+        params = dict(msgs=msgs, tokenizer=tokenizer, omni_input=True, **kw)
+        return repo.chat(**params
     elif filetype == "Audio":
         audio_np, sample_rate = librosa.load(input, sr=16000, mono=True)
         chunk_tensor = torch.from_numpy(audio_np).float().to(DEVICE)