Add support for vLLM and local models to GeminiVideoSummarization

.gitignore

@@ -63,4 +63,13 @@ cmake-build-release/
 *.pyc
 
 *.private
-venv
+venv
+*.mp4
+*.mov
+*.mkv
+*.png
+*.csv
+*.zip
+_*.py
+_*.txt
+_*.json

python/GeminiVideoSummarization/.dockerignore

@@ -0,0 +1,12 @@
+.venv/
+.vscode/
+**/__pycache__/
+*.egg-info/
+tests/*
+!tests/test_gemini.py
+!tests/data/
+tests/data/*
+!tests/data/NOTICE
+!tests/data/cat.mp4
+!tests/data/dog.mp4
+!tests/data/short.mp4

python/GeminiVideoSummarization/Dockerfile

@@ -0,0 +1,67 @@
+# syntax=docker/dockerfile:experimental
+
+#############################################################################
+# NOTICE                                                                    #
+#                                                                           #
+# This software (or technical data) was produced for the U.S. Government    #
+# under contract, and is subject to the Rights in Data-General Clause       #
+# 52.227-14, Alt. IV (DEC 2007).                                            #
+#                                                                           #
+# Copyright 2024 The MITRE Corporation. All Rights Reserved.                #
+#############################################################################
+
+#############################################################################
+# Copyright 2024 The MITRE Corporation                                      #
+#                                                                           #
+# Licensed under the Apache License, Version 2.0 (the "License");           #
+# you may not use this file except in compliance with the License.          #
+# You may obtain a copy of the License at                                   #
+#                                                                           #
+#    http://www.apache.org/licenses/LICENSE-2.0                             #
+#                                                                           #
+# Unless required by applicable law or agreed to in writing, software       #
+# distributed under the License is distributed on an "AS IS" BASIS,         #
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  #
+# See the License for the specific language governing permissions and       #
+# limitations under the License.                                            #
+#############################################################################
+
+ARG BUILD_REGISTRY
+ARG BUILD_TAG=latest
+FROM ${BUILD_REGISTRY}openmpf_python_executor_ssb:${BUILD_TAG} AS component
+
+RUN pip3 install --no-cache-dir \
+    tenacity \
+    numpy \
+    imageio-ffmpeg \
+    openai \
+    google-genai \
+    google-cloud-storage \
+    opencv-python-headless
+
+ARG RUN_TESTS=false
+
+RUN --mount=target=.,readwrite \
+    install-component.sh; \
+    if [ "${RUN_TESTS,,}" == true ]; then python tests/test_gemini.py; fi
+
+LABEL org.label-schema.license="Apache 2.0" \
+      org.label-schema.name="OpenMPF Gemini Video Summarization" \
+      org.label-schema.schema-version="1.1" \
+      org.label-schema.url="https://openmpf.github.io" \
+      org.label-schema.vcs-url="https://github.com/openmpf/openmpf-components" \
+      org.label-schema.vendor="MITRE"
+
+
+FROM component AS vllm-tests
+ARG RUN_TESTS=false
+ENV RUN_TESTS=${RUN_TESTS}
+
+WORKDIR /opt/gemini-tests
+COPY tests/test_gemini.py ./
+COPY tests/data/cat.mp4 data/cat.mp4
+
+ENTRYPOINT ["/bin/bash", "-lc", "if [ \"${RUN_TESTS,,}\" == \"true\" ]; then exec /opt/mpf/plugin-venv/bin/python3 -m unittest -v test_gemini.TestGemini.test_vllm_openai_compatible_api; fi"]
+
+
+FROM component AS final

python/GeminiVideoSummarization/Dockerfile.vllm

@@ -0,0 +1,115 @@
+# syntax=docker/dockerfile:1
+
+FROM vllm/vllm-openai:latest
+
+ENV HF_HOME=/root/.cache/huggingface
+ENV HF_HUB_OFFLINE=1
+
+ARG DEFAULT_VLLM_MODEL=google/gemma-4-12B-it
+ARG DEFAULT_VLLM_DRAFT_MODEL=
+
+RUN pip install --upgrade transformers vllm vllm[audio]
+
+
+# vLLM 0.22 knows how to turn gemma4_assistant into gemma4_mtp, but
+# current Transformers loads google/gemma-4-*-it-assistant as
+# gemma4_unified_assistant. Normalize that config type before vLLM's override.
+RUN python3 - <<'PY_SHIM'
+from pathlib import Path
+import site
+
+site_packages = Path(site.getsitepackages()[0])
+shim = site_packages / "gemma4_unified_assistant_mtp_shim.py"
+shim.write_text('''\
+from transformers.models.gemma4_unified.configuration_gemma4_unified import Gemma4UnifiedConfig
+from vllm.config.speculative import SpeculativeConfig
+
+def _alias_token_index(token_name):
+    index_name = f"{token_name}_index"
+    id_name = f"{token_name}_id"
+    if not hasattr(Gemma4UnifiedConfig, index_name):
+        setattr(
+            Gemma4UnifiedConfig,
+            index_name,
+            property(lambda self, attr=id_name: getattr(self, attr)),
+        )
+
+for _token_name in ("image_token", "video_token", "audio_token"):
+    _alias_token_index(_token_name)
+
+_original_hf_config_override = SpeculativeConfig.hf_config_override
+
+def _gemma4_unified_assistant_hf_config_override(hf_config):
+    if getattr(hf_config, "model_type", None) == "gemma4_unified_assistant":
+        hf_config.model_type = "gemma4_assistant"
+    return _original_hf_config_override(hf_config)
+
+SpeculativeConfig.hf_config_override = staticmethod(
+    _gemma4_unified_assistant_hf_config_override
+)
+''')
+(site_packages / "gemma4_unified_assistant_mtp_shim.pth").write_text(
+    "import gemma4_unified_assistant_mtp_shim\n"
+)
+PY_SHIM
+
+ENV VLLM_MODEL=${DEFAULT_VLLM_MODEL}
+ENV VLLM_DRAFT_MODEL=${DEFAULT_VLLM_DRAFT_MODEL}
+ENV VLLM_HOST=0.0.0.0
+ENV VLLM_PORT=8000
+ENV VLLM_TENSOR_PARALLEL_SIZE=1
+ENV VLLM_DRAFT_TENSOR_PARALLEL_SIZE=1
+ENV VLLM_DTYPE=auto
+ENV MPF_VLLM_ATTENTION_BACKEND=TRITON_ATTN
+# FP8 Marlin currently fails during Gemma 4 multimodal video profiling.
+ENV VLLM_QUANTIZATION=
+ENV VLLM_KV_CACHE_DTYPE=auto
+ENV VLLM_MAX_MODEL_LEN=131072
+ENV VLLM_MAX_NUM_BATCHED_TOKENS=8192
+ENV VLLM_MAX_NUM_SEQS=1
+ENV VLLM_NUM_SPECULATIVE_TOKENS=1
+ENV VLLM_ENFORCE_EAGER=1
+ENV VLLM_SPECULATIVE_ENFORCE_EAGER=1
+ENV MPF_LIMIT_MM_PER_PROMPT='{"video":1,"audio":1}'
+ENV VLLM_ENABLE_MTP=0
+
+EXPOSE 8000
+
+ENTRYPOINT []
+CMD bash -lc '\
+  set -euo pipefail && \
+  draft_model="${VLLM_DRAFT_MODEL:-${VLLM_MODEL}-assistant}" && \
+  echo "Serving ${VLLM_MODEL}" && \
+  set -- vllm serve "${VLLM_MODEL}" && \
+  set -- "$@" --host "${VLLM_HOST}" && \
+  set -- "$@" --port "${VLLM_PORT}" && \
+  set -- "$@" --tensor-parallel-size "${VLLM_TENSOR_PARALLEL_SIZE}" && \
+  set -- "$@" --dtype "${VLLM_DTYPE}" && \
+  if [ -n "${MPF_VLLM_ATTENTION_BACKEND}" ]; then \
+    set -- "$@" --attention-backend "${MPF_VLLM_ATTENTION_BACKEND}"; \
+  fi && \
+  set -- "$@" --kv_cache_dtype "${VLLM_KV_CACHE_DTYPE}" && \
+  set -- "$@" --max-model-len "${VLLM_MAX_MODEL_LEN}" && \
+  set -- "$@" --max-num-batched-tokens "${VLLM_MAX_NUM_BATCHED_TOKENS}" && \
+  if [ -n "${MPF_LIMIT_MM_PER_PROMPT}" ]; then \
+    set -- "$@" --limit-mm-per-prompt "${MPF_LIMIT_MM_PER_PROMPT}"; \
+  fi && \
+  if [ -n "${VLLM_MAX_NUM_SEQS}" ]; then \
+    set -- "$@" --max-num-seqs "${VLLM_MAX_NUM_SEQS}"; \
+  fi && \
+  if [ -n "${VLLM_QUANTIZATION}" ]; then \
+    set -- "$@" --quantization "${VLLM_QUANTIZATION}"; \
+  fi && \
+  if [ "${VLLM_ENFORCE_EAGER}" = "1" ]; then \
+    set -- "$@" --enforce-eager; \
+  fi && \
+  if [ "${VLLM_ENABLE_MTP}" = "1" ]; then \
+    echo "Using MTP draft model ${draft_model}"; \
+    speculative_enforce_eager=false; \
+    if [ "${VLLM_SPECULATIVE_ENFORCE_EAGER}" = "1" ]; then speculative_enforce_eager=true; fi; \
+    speculative_config=$(printf "{\"method\":\"mtp\",\"model\":\"%s\",\"num_speculative_tokens\":%s,\"draft_tensor_parallel_size\":%s,\"max_model_len\":%s,\"enforce_eager\":%s}" "${draft_model}" "${VLLM_NUM_SPECULATIVE_TOKENS}" "${VLLM_DRAFT_TENSOR_PARALLEL_SIZE}" "${VLLM_MAX_MODEL_LEN}" "${speculative_enforce_eager}"); \
+    set -- "$@" --speculative-config "${speculative_config}"; \
+  fi && \
+  unset VLLM_MODEL VLLM_DRAFT_MODEL VLLM_HOST VLLM_PORT VLLM_TENSOR_PARALLEL_SIZE VLLM_DRAFT_TENSOR_PARALLEL_SIZE VLLM_DTYPE VLLM_QUANTIZATION VLLM_KV_CACHE_DTYPE VLLM_MAX_MODEL_LEN VLLM_MAX_NUM_BATCHED_TOKENS VLLM_MAX_NUM_SEQS VLLM_NUM_SPECULATIVE_TOKENS VLLM_ENFORCE_EAGER VLLM_SPECULATIVE_ENFORCE_EAGER VLLM_ENABLE_MTP VLLM_BUILD_URL VLLM_IMAGE_TAG VLLM_BUILD_PIPELINE VLLM_BUILD_COMMIT && \
+  exec "$@"\
+'

python/GeminiVideoSummarization/Dockerfile.vllm-diffusiongemma

@@ -0,0 +1,110 @@
+# syntax=docker/dockerfile:1
+
+FROM vllm/vllm-openai:latest
+
+RUN pip install --upgrade transformers vllm vllm[audio]
+
+# Latest vLLM still falls back to its generic Transformers backend for
+# DiffusionGemma and requires decoder forward signatures to expose input_ids.
+RUN python3 - <<'PY_SHIM'
+from pathlib import Path
+import site
+
+site_packages = Path(site.getsitepackages()[0])
+shim = site_packages / "diffusiongemma_vllm_shim.py"
+shim.write_text('''\
+from transformers.models.diffusion_gemma import configuration_diffusion_gemma as _diffusion_gemma_config
+from transformers.models.diffusion_gemma import modeling_diffusion_gemma as _diffusion_gemma
+
+def _top_k_experts_alias(self):
+    return getattr(self, "top_k_experts", None)
+
+if not hasattr(_diffusion_gemma_config.DiffusionGemmaTextConfig, "num_experts_per_tok"):
+    _diffusion_gemma_config.DiffusionGemmaTextConfig.num_experts_per_tok = property(_top_k_experts_alias)
+if not hasattr(_diffusion_gemma_config.DiffusionGemmaTextConfig, "top_k"):
+    _diffusion_gemma_config.DiffusionGemmaTextConfig.top_k = property(_top_k_experts_alias)
+
+_original_forward = _diffusion_gemma.DiffusionGemmaDecoderModel.forward
+
+def _vllm_compatible_decoder_forward(
+    self,
+    input_ids=None,
+    inputs_embeds=None,
+    position_ids=None,
+    past_key_values=None,
+    self_conditioning_logits=None,
+    decoder_attention_mask=None,
+    decoder_position_ids=None,
+    decoder_input_ids=None,
+    **kwargs,
+):
+    if decoder_input_ids is None:
+        decoder_input_ids = input_ids
+    if decoder_position_ids is None:
+        decoder_position_ids = position_ids
+    return _original_forward(
+        self,
+        decoder_input_ids=decoder_input_ids,
+        past_key_values=past_key_values,
+        self_conditioning_logits=self_conditioning_logits,
+        decoder_attention_mask=decoder_attention_mask,
+        decoder_position_ids=decoder_position_ids,
+        **kwargs,
+    )
+
+_diffusion_gemma.DiffusionGemmaDecoderModel.forward = _vllm_compatible_decoder_forward
+''')
+(site_packages / "diffusiongemma_vllm_shim.pth").write_text("import diffusiongemma_vllm_shim\n")
+PY_SHIM
+
+ARG DEFAULT_MPF_VLLM_MODEL=google/diffusiongemma-26B-A4B-it
+
+ENV HF_HOME=/root/.cache/huggingface
+
+ENV MPF_VLLM_MODEL=${DEFAULT_MPF_VLLM_MODEL}
+ENV MPF_VLLM_HOST=0.0.0.0
+ENV MPF_VLLM_PORT=8001
+ENV MPF_VLLM_TENSOR_PARALLEL_SIZE=1
+ENV MPF_VLLM_DTYPE=auto
+ENV MPF_VLLM_LOAD_FORMAT=auto
+ENV MPF_VLLM_QUANTIZATION=fp8
+ENV MPF_VLLM_KV_CACHE_DTYPE=auto
+ENV MPF_VLLM_MAX_MODEL_LEN=8192
+ENV MPF_VLLM_MAX_NUM_BATCHED_TOKENS=2048
+ENV MPF_VLLM_MAX_NUM_SEQS=1
+ENV MPF_VLLM_GPU_MEMORY_UTILIZATION=0.90
+ENV MPF_VLLM_CPU_OFFLOAD_GB=20
+ENV MPF_VLLM_ENFORCE_EAGER=0
+ENV MPF_LIMIT_MM_PER_PROMPT=
+
+EXPOSE 8001
+
+ENTRYPOINT []
+CMD bash -lc '\
+  set -euo pipefail && \
+  unset VLLM_BUILD_URL VLLM_IMAGE_TAG VLLM_BUILD_PIPELINE VLLM_BUILD_COMMIT && \
+  echo "Serving ${MPF_VLLM_MODEL}" && \
+  set -- vllm serve "${MPF_VLLM_MODEL}" && \
+  set -- "$@" --host "${MPF_VLLM_HOST}" && \
+  set -- "$@" --port "${MPF_VLLM_PORT}" && \
+  set -- "$@" --tensor-parallel-size "${MPF_VLLM_TENSOR_PARALLEL_SIZE}" && \
+  set -- "$@" --dtype "${MPF_VLLM_DTYPE}" && \
+  set -- "$@" --load-format "${MPF_VLLM_LOAD_FORMAT}" && \
+  set -- "$@" --kv-cache-dtype "${MPF_VLLM_KV_CACHE_DTYPE}" && \
+  set -- "$@" --max-model-len "${MPF_VLLM_MAX_MODEL_LEN}" && \
+  set -- "$@" --max-num-batched-tokens "${MPF_VLLM_MAX_NUM_BATCHED_TOKENS}" && \
+  set -- "$@" --max-num-seqs "${MPF_VLLM_MAX_NUM_SEQS}" && \
+  set -- "$@" --gpu-memory-utilization "${MPF_VLLM_GPU_MEMORY_UTILIZATION}" && \
+  set -- "$@" --cpu-offload-gb "${MPF_VLLM_CPU_OFFLOAD_GB}" && \
+  if [ -n "${MPF_LIMIT_MM_PER_PROMPT}" ]; then \
+    set -- "$@" --limit-mm-per-prompt "${MPF_LIMIT_MM_PER_PROMPT}"; \
+  fi && \
+  if [ -n "${MPF_VLLM_QUANTIZATION}" ]; then \
+    set -- "$@" --quantization "${MPF_VLLM_QUANTIZATION}"; \
+  fi && \
+  if [ "${MPF_VLLM_ENFORCE_EAGER}" = "1" ]; then \
+    set -- "$@" --enforce-eager; \
+  fi && \
+  printf "vLLM command:" && printf " %q" "$@" && printf "\\n" && \
+  exec "$@"\
+'