Add Gemma3 support for videos

Author: tomasruizt

llmlib/llmlib/huggingface_inference.py

@@ -1,3 +1,4 @@
+import math
 import os
 import base64
 import io
@@ -7,13 +8,19 @@
 import PIL
 from enum import StrEnum
 from .base_llm import LLM, Message, validate_only_first_message_has_files
+import cv2
+from PIL import Image
+from logging import getLogger
+
+
+logger = getLogger(__name__)
 
 
 def get_image_as_base64(image_bytes: bytes):
     return base64.b64encode(image_bytes).decode("utf-8")
 
 
-def convert_message_to_hf_format(message: Message) -> dict:
+def convert_message_to_hf_format(message: Message, max_n_frames_per_video: int) -> dict:
     """Convert a Message to HuggingFace chat format."""
     content = []
 
@@ -25,20 +32,77 @@ def convert_message_to_hf_format(message: Message) -> dict:
     if message.img is not None:
         content.append(extract_content_piece(message.img))
 
-    # Add multiple images from files if present
+    if message.video is not None:
+        imgs: list = video_to_imgs(message.video, max_n_frames_per_video)
+        for frame in imgs:
+            content.append(extract_content_piece(frame))
+
+    # Add multiple images from files (img or video) if present
     if message.files is not None:
         for file_path in message.files:
             if is_img(file_path):
                 content.append(extract_content_piece(file_path))
+            elif is_video(file_path):
+                imgs: list = video_to_imgs(file_path, max_n_frames_per_video)
+                for frame in imgs:
+                    content.append(extract_content_piece(frame))
+            else:
+                raise ValueError(f"Unsupported file type: {file_path}")
 
     return {"role": message.role, "content": content}
 
 
+def video_to_imgs(video_path: Path, max_n_frames: int) -> list[PIL.Image.Image]:
+    assert isinstance(video_path, Path), video_path
+    """From https://github.com/agustoslu/simple-inference-benchmark/blob/5cec55787d34af65f0d11efc429c3d4de92f051a/utils.py#L79"""
+    cap = cv2.VideoCapture(str(video_path))
+    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    fps = int(cap.get(cv2.CAP_PROP_FPS))
+
+    frame_indices = compute_frame_indices(
+        vid_n_frames=total_frames, vid_fps=fps, max_n_frames=max_n_frames
+    )
+
+    frames = []
+    for frame_idx in frame_indices:
+        cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
+        success, frame = cap.read()
+        if success:
+            # Convert BGR (the default format for OpenCV) to RGB
+            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            frames.append(Image.fromarray(frame_rgb))
+
+    cap.release()
+    logger.info(f"Extracted {len(frames)} frames from video {video_path}")
+    return frames
+
+
+def compute_frame_indices(vid_n_frames: int, vid_fps: float, max_n_frames: int):
+    """
+    From https://github.com/agustoslu/simple-inference-benchmark/blob/5cec55787d34af65f0d11efc429c3d4de92f051a/utils.py#L164
+    This function will return the frames starting at 0 every second.
+    Unless that number exceeds max_n_frames, in which case it will return max_n_frames frames evenly spaced out, starting at 0.
+    """
+    assert isinstance(vid_n_frames, int), vid_n_frames
+    assert isinstance(max_n_frames, int), max_n_frames
+    vid_fps = int(vid_fps)
+    fps_n_frames = math.ceil(vid_n_frames / vid_fps)
+    if fps_n_frames <= max_n_frames:
+        return list(range(0, vid_n_frames - 1, vid_fps))
+    else:
+        return list(range(0, vid_n_frames - 1, vid_n_frames // max_n_frames))
+
+
 def is_img(file_path: str | Path) -> bool:
     permitted = (".png", ".jpg", ".jpeg")
     return str(file_path).lower().endswith(permitted)
 
 
+def is_video(file_path: str | Path) -> bool:
+    permitted = (".mp4",)
+    return str(file_path).lower().endswith(permitted)
+
+
 def extract_content_piece(img: PIL.Image.Image | str | Path) -> dict:
     image_bytes = extract_bytes(img)
     content_piece = {
@@ -58,7 +122,7 @@ def extract_bytes(img: PIL.Image.Image | str | Path) -> bytes:
     elif isinstance(img, PIL.Image.Image):
         # Handle PIL Image
         img_byte_arr = io.BytesIO()
-        img.save(img_byte_arr, format=img.format or "PNG")
+        img.save(img_byte_arr, format="jpeg")
         return img_byte_arr.getvalue()
     else:
         raise ValueError(f"Unsupported image type: {type(img)}")
@@ -75,6 +139,7 @@ class HuggingFaceVLM(LLM):
     model_id: HuggingFaceVLMs
     max_new_tokens: int = 1000
     requires_gpu_exclusively: bool = False
+    max_n_frames_per_video: int = 200
 
     # Available model IDs
     model_ids = list(HuggingFaceVLMs)
@@ -92,13 +157,20 @@ def __post_init__(self):
     def complete_msgs(self, msgs: list[Message]) -> str:
         """Complete a conversation with the model."""
         validate_only_first_message_has_files(msgs)
-        hf_messages = [convert_message_to_hf_format(msg) for msg in msgs]
+        hf_messages = [
+            convert_message_to_hf_format(
+                msg, max_n_frames_per_video=self.max_n_frames_per_video
+            )
+            for msg in msgs
+        ]
 
+        logger.info("Calling HuggingFace API...")
         completion = self.client.chat.completions.create(
             model=self.model_id,
             messages=hf_messages,
             max_tokens=self.max_new_tokens,
         )
+        logger.info("Token usage: %s", dict(completion.usage))
 
         return completion.choices[0].message.content
 

test-files/.gitignore

@@ -3,3 +3,4 @@
 *.mp3
 *.mp4
 *.flac
+generated_*

tests/helpers.py

@@ -1,7 +1,10 @@
+import base64
 import os
 from pathlib import Path
 import PIL
+import cv2
 from llmlib.base_llm import LLM, Message
+import numpy as np
 import pytest
 
 
@@ -75,8 +78,10 @@ def mona_lisa_message() -> Message:
     return msg
 
 
-def pyramid_message() -> Message:
+def pyramid_message(load_img: bool = False) -> Message:
     img = file_for_test("pyramid.jpg")
+    if load_img:
+        img = PIL.Image.open(img)
     msg = Message(role="user", msg="What is in the image?", img=img, img_name="")
     return msg
 
@@ -131,8 +136,7 @@ def assert_model_supports_multiturn(model: LLM):
 
 
 def assert_model_supports_multiturn_with_6min_video(model: LLM):
-    video = file_for_test("tasting travel - rome italy.mp4")
-    convo = [Message(role="user", msg="What country are they visiting?", video=video)]
+    convo = [video_message()]
     answer1 = model.complete_msgs(convo)
     assert "italy" in answer1.lower(), answer1
 
@@ -149,6 +153,11 @@ def assert_model_supports_multiturn_with_6min_video(model: LLM):
     assert "jesus" in answer3.lower(), answer3
 
 
+def video_message() -> Message:
+    video = file_for_test("tasting travel - rome italy.mp4")
+    return Message(role="user", msg="What country are they visiting?", video=video)
+
+
 def assert_model_supports_multiturn_with_multiple_imgs(model: LLM):
     files = [file_for_test("forest.jpg"), file_for_test("fish.jpg")]
     msg = Message(
@@ -164,3 +173,14 @@ def assert_model_supports_multiturn_with_multiple_imgs(model: LLM):
     answer2 = model.complete_msgs(convo).lower()
     possible_answers = ["biodiversity", "ecosystem", "habitat"]
     assert any(answer in answer2 for answer in possible_answers), answer2
+
+
+def decode_base64_to_array(base64_str: str) -> np.ndarray:
+    """Decode base64 string to OpenCV image (numpy array)"""
+    # Remove data URL prefix if present
+    if "base64," in base64_str:
+        base64_str = base64_str.split("base64,")[1]
+    image_data = base64.b64decode(base64_str)
+    np_array = np.frombuffer(image_data, np.uint8)
+    image = cv2.imdecode(np_array, cv2.IMREAD_COLOR)
+    return image

tests/test_huggingface_vlm.py

@@ -1,11 +1,18 @@
+import cv2
+from llmlib.huggingface_inference import convert_message_to_hf_format
 import pytest
 from llmlib.huggingface_inference import HuggingFaceVLM, HuggingFaceVLMs
 from .helpers import (
     assert_model_recognizes_pyramid_in_image,
+    assert_model_supports_multiturn_with_6min_video,
+    decode_base64_to_array,
+    file_for_test,
     is_ci,
     assert_model_knows_capital_of_france,
     assert_model_supports_multiturn,
     assert_model_supports_multiturn_with_multiple_imgs,
+    pyramid_message,
+    video_message,
 )
 
 
@@ -45,3 +52,31 @@ def test_huggingface_vlm_multi_turn_text_conversation(gemma3):
 @pytest.mark.skipif(condition=is_ci(), reason="Avoid costs")
 def test_huggingface_vlm_multi_turn_with_images(gemma3):
     assert_model_supports_multiturn_with_multiple_imgs(gemma3)
+
+
+@pytest.mark.skipif(condition=is_ci(), reason="Files are not available on CI")
+def test_huggingface_vlm_multi_turn_with_6min_video(gemma3):
+    assert_model_supports_multiturn_with_6min_video(gemma3)
+
+
+@pytest.mark.skipif(condition=is_ci(), reason="Files are not available on CI")
+def test_convert_to_huggingface_format():
+    img_msg1 = pyramid_message(load_img=True)
+    img_msg2 = pyramid_message(load_img=False)
+    max_n_frames_per_video = 200
+    b64_enc1 = convert_message_to_hf_format(img_msg1, max_n_frames_per_video)[
+        "content"
+    ][1]["image_url"]["url"]
+    b64_enc2 = convert_message_to_hf_format(img_msg2, max_n_frames_per_video)[
+        "content"
+    ][1]["image_url"]["url"]
+    # assert b64_enc1 == b64_enc2
+    array1 = decode_base64_to_array(base64_str=b64_enc1)
+    array2 = decode_base64_to_array(base64_str=b64_enc2)
+    # asserting imgs are the same fails, but you can visually inspect them
+    cv2.imwrite(file_for_test("generated_pyramid_1.jpeg"), array1)
+    cv2.imwrite(file_for_test("generated_pyramid_2.jpeg"), array2)
+
+    msg = video_message()
+    hf_msg = convert_message_to_hf_format(msg, max_n_frames_per_video)
+    assert len(hf_msg["content"]) > 10