Update app.py

app.py

@@ -7,6 +7,7 @@ import math
 import librosa
 from PIL import Image, ImageSequence
 from decord import VideoReader, cpu
+from moviepy.editor import VideoFileClip
 from transformers import AutoModel, AutoTokenizer, AutoProcessor
 
 # Variables
@@ -51,13 +52,15 @@ def uniform_sample(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)
+    clip = VideoFileClip(path)
+    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
+        clip.audio.write_audiofile(tmp.name, fps=sr, codec="pcm_s16le", verbose=False, logger=None)
+        audio_np, _ = librosa.load(tmp.name, sr=sr, mono=True)
+    total_units = math.ceil(clip.duration / seconds_per_unit)
     content = []
     for i in range(total_units):
-        frame = Image.fromarray(vr[int(i * fps * seconds_per_unit)].asnumpy().astype("uint8"))
+        t = min(i * seconds_per_unit, clip.duration - 1e-3)
+        frame = Image.fromarray(clip.get_frame(t).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