fix(customizer): fix dockerfile organization for customizer, add AV opencv build with ffmpeg 8.1.2+

docker-bake.hcl

@@ -52,6 +52,10 @@ variable "MAMBA_SSM_WHEEL_CONTEXT" {
   default = ""
 }
 
+variable "FFMPEG_VLM_WHEEL_CONTEXT" {
+  default = ""
+}
+
 variable "DISTROLESS_BASE" {
   default = "nvcr.io/nvidia/distroless/python:3.11-v4.0.8"
 }
@@ -75,12 +79,12 @@ variable "BASE_TAG_PYTHON" {
 
 # Pin for nmp-automodel-base.
 variable "BASE_TAG_AUTOMODEL" {
-  default = "f8239353044d71cbd53e209b60c0600ead484b58"
+  default = "f0756dd64eaf2ddb9c5c962e18216b2e70ba4b64"
 }
 
 # The tag for base images if needed
 variable "WHEELS_TAG" {
-  default = "65632527b258367fb8855ef431f596abb8538577"
+  default = "f0756dd64eaf2ddb9c5c962e18216b2e70ba4b64"
 }
 
 variable "BAKE_CACHE_SOURCE_BRANCH" {
@@ -138,6 +142,11 @@ function "get_mamba_ssm_wheel_image" {
   result = "${WHEELS_REGISTRY}/mamba-ssm-wheel:${WHEELS_TAG}"
 }
 
+function "get_ffmpeg_vlm_wheel_image" {
+  params = []
+  result = "${WHEELS_REGISTRY}/ffmpeg-vlm-wheel:${WHEELS_TAG}"
+}
+
 function "get_arch_tag" {
   params = []
   result = BUILD_ARCH == "linux/arm64" ? "linux-arm64" : "linux-amd64"
@@ -175,6 +184,11 @@ function "mamba_ssm_wheel_context" {
   result = notequal(MAMBA_SSM_WHEEL_CONTEXT, "") ? MAMBA_SSM_WHEEL_CONTEXT : notequal(USE_LOCAL_WHEELS, "") ? "target:mamba-ssm-wheel" : "docker-image://${get_mamba_ssm_wheel_image()}"
 }
 
+function "ffmpeg_vlm_wheel_context" {
+  params = []
+  result = notequal(FFMPEG_VLM_WHEEL_CONTEXT, "") ? FFMPEG_VLM_WHEEL_CONTEXT : notequal(USE_LOCAL_WHEELS, "") ? "target:ffmpeg-vlm-wheel" : "docker-image://${get_ffmpeg_vlm_wheel_image()}"
+}
+
 function "wheel_tags" {
   params = [name]
   result = [
@@ -293,6 +307,7 @@ group "nmp-automodel-gpu-wheels" {
   targets = [
     "causal-conv1d-wheel",
     "mamba-ssm-wheel",
+    "ffmpeg-vlm-wheel",
   ]
 }
 
@@ -500,15 +515,14 @@ target "nmp-cpu-tasks-docker" {
   platforms  = get_platforms()
 }
 
-# Mamba Wheel Builders
-# Builds Python wheels for mamba-ssm and causal-conv1d in parallel.
-# Both only ship source distributions on PyPI; this pre-builds them for
-# amd64 and arm64. The wheels live at /wheels/*.whl inside the image.
+# Python wheel builders (causal-conv1d, mamba-ssm, av, opencv-python-headless).
+# CUDA extensions only ship source on PyPI; av/opencv bundle FFmpeg. Pre-built for
+# amd64 and arm64. Wheels live at /wheels/*.whl inside each image.
 
 target "causal-conv1d-wheel" {
   target     = "causal-conv1d-wheel"
   context    = "."
-  dockerfile = "docker/base/Dockerfile.mamba-wheel"
+  dockerfile = "docker/base/Dockerfile.python-wheels"
   cache-to   = maybe_registry_cache_to("causal-conv1d-wheel")
   cache-from = maybe_registry_cache_from("causal-conv1d-wheel")
   tags       = wheel_tags("causal-conv1d-wheel")
@@ -523,7 +537,7 @@ target "causal-conv1d-wheel" {
 target "mamba-ssm-wheel" {
   target     = "mamba-ssm-wheel"
   context    = "."
-  dockerfile = "docker/base/Dockerfile.mamba-wheel"
+  dockerfile = "docker/base/Dockerfile.python-wheels"
   cache-to   = maybe_registry_cache_to("mamba-ssm-wheel")
   cache-from = maybe_registry_cache_from("mamba-ssm-wheel")
   tags       = wheel_tags("mamba-ssm-wheel")
@@ -536,6 +550,17 @@ target "mamba-ssm-wheel" {
   platforms = get_platforms()
 }
 
+target "ffmpeg-vlm-wheel" {
+  target     = "ffmpeg-vlm-wheel"
+  context    = "."
+  dockerfile = "docker/base/Dockerfile.python-wheels"
+  cache-to   = maybe_registry_cache_to("ffmpeg-vlm-wheel")
+  cache-from = maybe_registry_cache_from("ffmpeg-vlm-wheel")
+  tags       = wheel_tags("ffmpeg-vlm-wheel")
+  output     = image_output()
+  platforms  = get_platforms()
+}
+
 
 target "safe-synthesizer-tasks-docker" {
   target     = "runtime"
@@ -698,7 +723,7 @@ target "automodel-platform-workspace" {
 target "nmp-automodel-base-builder" {
   target          = "nmp-automodel-base"
   context         = "."
-  dockerfile      = "docker/Dockerfile.nmp-automodel-base"
+  dockerfile      = "docker/automodel/Dockerfile.nmp-automodel-base"
   no-cache-filter = ["automodel-clone"]
   cache-to        = maybe_registry_cache_to("nmp-automodel-base")
   cache-from      = maybe_registry_cache_from("nmp-automodel-base")
@@ -707,14 +732,15 @@ target "nmp-automodel-base-builder" {
   contexts = {
     causal-conv1d-wheel-image = causal_conv1d_wheel_context()
     mamba-ssm-wheel-image     = mamba_ssm_wheel_context()
+    ffmpeg-vlm-wheel-image    = ffmpeg_vlm_wheel_context()
   }
   platforms = get_platforms()
 }
 
 target "nmp-automodel-tasks-docker" {
   target     = "runtime"
   context    = "."
-  dockerfile = "docker/Dockerfile.nmp-automodel-tasks"
+  dockerfile = "docker/automodel/Dockerfile.nmp-automodel-tasks"
   contexts = {
     platform-workspace = "target:automodel-platform-workspace"
     nmp-automodel-base = automodel_base_context()
@@ -729,7 +755,7 @@ target "nmp-automodel-tasks-docker" {
 target "nmp-automodel-training-docker" {
   target     = "runtime"
   context    = "."
-  dockerfile = "docker/Dockerfile.nmp-automodel-training"
+  dockerfile = "docker/automodel/Dockerfile.nmp-automodel-training"
   contexts = {
     platform-workspace = "target:automodel-platform-workspace"
     nmp-automodel-base = automodel_base_context()
@@ -744,7 +770,7 @@ target "nmp-automodel-training-docker" {
 target "nmp-automodel-tasks-smoke-test" {
   target     = "smoke-test"
   context    = "."
-  dockerfile = "docker/Dockerfile.nmp-automodel-tasks"
+  dockerfile = "docker/automodel/Dockerfile.nmp-automodel-tasks"
   contexts = {
     platform-workspace = "target:automodel-platform-workspace"
     nmp-automodel-base = automodel_base_context()
@@ -760,7 +786,7 @@ target "nmp-automodel-tasks-smoke-test" {
 target "nmp-automodel-training-smoke-test" {
   target     = "smoke-test"
   context    = "."
-  dockerfile = "docker/Dockerfile.nmp-automodel-training"
+  dockerfile = "docker/automodel/Dockerfile.nmp-automodel-training"
   contexts = {
     platform-workspace = "target:automodel-platform-workspace"
     nmp-automodel-base = automodel_base_context()

docker/Dockerfile.nmp-unsloth-training

@@ -13,10 +13,10 @@
 #      blocks uv from installing/upgrading torch into the venv so the NGC
 #      base's PyTorch + CUDA remain the runtime stack.
 #   1b. bitsandbytes — compiled from source against NGC CUDA 13.2 (same pattern
-#       as docker/Dockerfile.nmp-automodel-base). PyPI wheels only ship through
+#       as docker/automodel/Dockerfile.nmp-automodel-base). PyPI wheels only ship through
 #       cuda130; source build replaces the wheel from step 1.
 #   1c. mamba-ssm + causal-conv1d — prebuilt cu13.2 / cp312 wheels (shared with
-#       docker/Dockerfile.nmp-automodel-base).
+#       docker/automodel/Dockerfile.nmp-automodel-base).
 #   1d. flash-attn — optional for unsloth and currently NOT installed (see the
 #       commented TODO below). Without it Unsloth falls back when xformers is
 #       also missing (common on newer CUDA stacks), logging "FA2 = False /
@@ -28,14 +28,14 @@
 # Default tag:    `local` (override via BAKE_TAG at build time).
 
 # NGC PyTorch base. 26.05-py3 ships PyTorch 2.12 + CUDA 13.2 + Python 3.12 (aligned with
-# docker/Dockerfile.nmp-automodel-base). Override at build time:
+# docker/automodel/Dockerfile.nmp-automodel-base). Override at build time:
 # --set nmp-unsloth-training.args.PYTORCH_BASE=...
 ARG PYTORCH_BASE=nvcr.io/nvidia/pytorch:26.05-py3
 
 # Prebuilt CUDA-extension wheels (mamba-ssm + causal-conv1d), shared with
-# docker/Dockerfile.nmp-automodel-base. The bake `nmp-unsloth-training` target
+# docker/automodel/Dockerfile.nmp-automodel-base. The bake `nmp-unsloth-training` target
 # resolves these contexts to the causal-conv1d-wheel / mamba-ssm-wheel images
-# (see docker-bake.hcl + docker/base/Dockerfile.mamba-wheel). Local builds need
+# (see docker-bake.hcl + docker/base/Dockerfile.python-wheels). Local builds need
 # `USE_LOCAL_WHEELS=1` (build the wheel targets) or a reachable WHEELS_REGISTRY.
 FROM causal-conv1d-wheel-image AS causal-conv1d-wheel-src
 FROM mamba-ssm-wheel-image AS mamba-ssm-wheel-src

docker/Dockerfile.nmp-automodel-base → docker/automodel/Dockerfile.nmp-automodel-base

@@ -11,6 +11,7 @@ ARG PYTORCH_BASE=nvcr.io/nvidia/pytorch:26.05-py3
 
 FROM causal-conv1d-wheel-image AS causal-conv1d-wheel-src
 FROM mamba-ssm-wheel-image AS mamba-ssm-wheel-src
+FROM ffmpeg-vlm-wheel-image AS ffmpeg-vlm-wheel-src
 
 FROM alpine/git AS automodel-clone
 ARG AUTOMODEL_COMMIT
@@ -111,36 +112,20 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     "grpcio>=1.81.1,<2" \
     "wandb>=0.28.0,<1"
 
-# <ffmpeg-fix>
-# Bundled FFmpeg in PyAV / opencv wheels is still below nspect 8.1.2 (av 17.1.0 → 8.1.1,
-# opencv-python-headless 4.13.0.92 → 8.0.1). Re-enable these pins once upstream wheels
-# ship FFmpeg >= 8.1.2, then drop the rm -rf cleanup step below.
-#
-# Expected to break without av / opencv (until pins are restored):
-#   - Qwen3-Omni / Qwen2.5-Omni VLM (qwen_omni_utils imports av; qwen3_omni_collate_fn)
-#   - Omni use_audio_in_video (av.open to probe video audio tracks)
-#   - Mistral / Pixtral / Mistral3 image inputs (mistral_common cv2.resize in image tokenizer)
-#   - import albumentations (cv2 required; listed in vlm extra but unused by Automodel today)
-# Still expected to work:
-#   - Qwen VL fine-tune (FORCE_QWENVL_VIDEO_READER=torchcodec; no av import in qwen_vl_utils)
-#   - Qwen3-Omni ASR-only collate (audio/collate_fns avoids qwen_omni_utils)
-#   - Mistral / Pixtral text-only tokenization
-# RUN --mount=type=cache,target=/root/.cache/uv \
-#     uv pip install \
-#     "av>=17.1.0,<18" \
-#     "opencv-python-headless>=4.13.0.92,<4.14"
-
-# Remove PyAV / opencv from the venv so nspect does not scan bundled libav* in av.libs /
-# opencv_python_headless.libs. vlm sync still pulls them transitively (qwen-omni-utils,
-# albumentations, mistral-common[opencv]); see breakage notes above.
-RUN rm -rf \
+# Replace PyPI av/opencv (bundled FFmpeg < 8.1.2) with custom wheels from ffmpeg-vlm-wheel.
+# vlm sync pulls PyPI copies transitively; drop them before installing our manylinux wheels.
+# PyAV 17.x uses cp311-abi3 (stable ABI) — compatible with cp312, not cp312-cp312 tagged.
+RUN --mount=from=ffmpeg-vlm-wheel-src,target=/tmp/ffmpeg-vlm-wheel-src,readonly \
+    rm -rf \
     /opt/venv/lib/python*/site-packages/av \
     /opt/venv/lib/python*/site-packages/av-*.dist-info \
     /opt/venv/lib/python*/site-packages/av.libs \
     /opt/venv/lib/python*/site-packages/cv2 \
     /opt/venv/lib/python*/site-packages/opencv_python_headless.libs \
-    /opt/venv/lib/python*/site-packages/opencv_python_headless-*.dist-info
-# </ffmpeg-fix>
+    /opt/venv/lib/python*/site-packages/opencv_python_headless-*.dist-info && \
+    uv pip install --no-cache-dir --no-deps \
+    /tmp/ffmpeg-vlm-wheel-src/wheels/av-*-abi3-manylinux_2_28_*.whl \
+    /tmp/ffmpeg-vlm-wheel-src/wheels/opencv_python_headless-*cp312*.whl
 
 # Published base image (same filesystem as builder).
 FROM ${PYTORCH_BASE} AS nmp-automodel-base

docker/automodel/README.md

@@ -1,12 +1,14 @@
 # nmp-automodel container images
 
+All Automodel Docker build files live under **`docker/automodel/`** (wheel vendor script: `docker/base/build-ffmpeg-vendor.sh`).
+
 Three images for the **nmp-automodel** customization backend. Published as flat repo names under **`my-registry/nemo-platform-dev/nmp-automodel-*`** (no nested `nmp/...` path — some registries reject that on push).
 
 | Image | Dockerfile | Role |
 |-------|------------|------|
-| `nmp-automodel-base` | `Dockerfile.nmp-automodel-base` | PyTorch 26.05 + Automodel + `mamba-ssm` / `causal-conv1d` wheels |
-| `nmp-automodel-tasks` | `Dockerfile.nmp-automodel-tasks` | Platform task glue (`file_io`, `model_entity`, `model_spec`); GPU-capable base |
-| `nmp-automodel-training` | `Dockerfile.nmp-automodel-training` | Training step (`nmp.automodel.tasks.training`) |
+| `nmp-automodel-base` | `docker/automodel/Dockerfile.nmp-automodel-base` | PyTorch 26.05 + Automodel + `mamba-ssm` / `causal-conv1d` wheels |
+| `nmp-automodel-tasks` | `docker/automodel/Dockerfile.nmp-automodel-tasks` | Platform task glue (`file_io`, `model_entity`, `model_spec`); GPU-capable base |
+| `nmp-automodel-training` | `docker/automodel/Dockerfile.nmp-automodel-training` | Training step (`nmp.automodel.tasks.training`) |
 
 Full references (default tag `local`):
 

docker/automodel/pyproject.workspace.toml

@@ -8,6 +8,8 @@ requires-python = ">=3.11,<3.14"
 
 [tool.uv]
 required-version = ">=0.9.14,<0.10.0"
+# greenlet 3.5.x ships x86_64-only wheels today; cap until aarch64 wheels land.
+constraint-dependencies = ["greenlet>=3.0.0,<3.5"]
 
 [tool.uv.workspace]
 members = [

docker/base/Dockerfile.mamba-wheel → docker/base/Dockerfile.python-wheels

@@ -1,21 +1,24 @@
 # syntax=docker/dockerfile:1
 #######
-# Mamba Wheel Builder
+# Python Wheel Builder
 #
 # Builds Python wheels for:
 #   - causal-conv1d (CUDA extension required by mamba-ssm)
 #   - mamba-ssm     (selective state space model)
+#   - av + opencv-python-headless (FFmpeg 8.1.2 bundled wheels for nspect / VLM)
 #
-# Both only ship source distributions on PyPI and require nvcc to compile.
-# The two builds are independent stages so BuildKit runs them in parallel.
-# Each image stores its wheel at /wheels/*.whl.
+# causal-conv1d / mamba-ssm only ship source on PyPI and require nvcc.
+# av / opencv wheels bundle FFmpeg; upstream PyPI wheels lag nspect 8.1.2.
 # Build via the nmp-automodel-gpu-wheels bake group.
 #
 # Build args:
 #   CAUSAL_CONV1D_VERSION  - git tag to build              (default: v1.5.3)
 #   MAMBA_22_COMMIT        - git SHA or tag to build
 #   CUDA_VERSION           - CUDA devel image version      (default: 12.8.1)
 #   TORCH_CUDA_ARCH_LIST   - semicolon-separated SM targets (default: "8.0;8.6;9.0")
+#   FFMPEG_VERSION         - FFmpeg for av/opencv wheels   (default: 8.1.2)
+#   PYAV_VERSION           - PyAV git tag                  (default: v17.1.0)
+#   OPENCV_PYTHON_GIT_TAG   - opencv-python git tag (92 = PyPI 4.13.0.92)
 #######
 
 ARG CUDA_VERSION=12.8.1
@@ -271,3 +274,85 @@ COPY --from=mamba-ssm-25-wheel-builder-py312 /wheels /wheels
 COPY --from=mamba-ssm-23-wheel-builder-py312 /wheels /wheels
 COPY --from=mamba-ssm-23-wheel-builder-py312-cu13.1.1 /wheels /wheels/cu13.1.1
 COPY --from=mamba-ssm-23-wheel-builder-py312-cu13.2 /wheels /wheels/cu13.2
+
+# =============================================================================
+# FFmpeg / VLM wheels (manylinux_2_28, cp312) — av + opencv-python-headless
+# =============================================================================
+FROM quay.io/pypa/manylinux_2_28_x86_64:latest AS ffmpeg-vendor-base-py312-amd64
+FROM quay.io/pypa/manylinux_2_28_aarch64:latest AS ffmpeg-vendor-base-py312-arm64
+
+ARG TARGETARCH
+FROM ffmpeg-vendor-base-py312-${TARGETARCH} AS ffmpeg-vendor-base-py312
+
+ARG FFMPEG_VERSION=8.1.2
+
+ENV FFMPEG_VERSION=${FFMPEG_VERSION} \
+    FFMPEG_PREFIX=/ffmpeg_build \
+    PKG_CONFIG_PATH=/ffmpeg_build/lib/pkgconfig:/usr/local/lib/pkgconfig \
+    LD_LIBRARY_PATH=/ffmpeg_build/lib:${LD_LIBRARY_PATH:-}
+
+RUN yum install -y \
+    curl \
+    git \
+    nasm \
+    yasm \
+    openssl-devel \
+    pkgconfig \
+    && yum clean all
+
+COPY docker/base/build-ffmpeg-vendor.sh /usr/local/bin/build-ffmpeg-vendor.sh
+RUN chmod +x /usr/local/bin/build-ffmpeg-vendor.sh && \
+    build-ffmpeg-vendor.sh
+
+# =============================================================================
+# av wheel — PyAV against shared FFmpeg 8.1.2
+# =============================================================================
+FROM ffmpeg-vendor-base-py312 AS av-wheel-builder-py312
+
+ARG PYAV_VERSION=v17.1.0
+
+ENV PKG_CONFIG_PATH=/ffmpeg_build/lib/pkgconfig \
+    LD_LIBRARY_PATH=/ffmpeg_build/lib:${LD_LIBRARY_PATH:-}
+
+RUN mkdir -p /wheels /tmp/wheelhouse && \
+    /opt/python/cp312-cp312/bin/pip install -U "pip<26" "cython>=3.1,<4" "numpy>=2.0" auditwheel && \
+    git clone --depth 1 --branch ${PYAV_VERSION} https://github.com/PyAV-Org/PyAV.git /src/PyAV && \
+    cd /src/PyAV && \
+    /opt/python/cp312-cp312/bin/pip wheel . --no-deps -w /tmp/wheelhouse && \
+    auditwheel repair /tmp/wheelhouse/av-*.whl -w /wheels && \
+    rm -rf /src/PyAV /tmp/wheelhouse
+
+# =============================================================================
+# opencv-python-headless wheel — opencv-python 4.13 against shared FFmpeg 8.1.2
+# =============================================================================
+FROM ffmpeg-vendor-base-py312 AS opencv-headless-wheel-builder-py312
+
+ARG OPENCV_PYTHON_GIT_TAG=92
+
+ENV CI_BUILD=1 \
+    ENABLE_HEADLESS=1 \
+    PKG_CONFIG_PATH=/ffmpeg_build/lib/pkgconfig:/usr/local/lib/pkgconfig \
+    LD_LIBRARY_PATH=/ffmpeg_build/lib:${LD_LIBRARY_PATH:-} \
+    LDFLAGS=-L/ffmpeg_build/lib
+
+RUN yum install -y \
+    zlib-devel \
+    libjpeg-turbo-devel \
+    && yum clean all
+
+RUN mkdir -p /wheels /tmp/wheelhouse && \
+    /opt/python/cp312-cp312/bin/pip install -U "pip<26" "numpy>=2.0" auditwheel && \
+    git clone --depth 1 --branch ${OPENCV_PYTHON_GIT_TAG} https://github.com/opencv/opencv-python.git /src/opencv-python && \
+    cd /src/opencv-python && \
+    git submodule update --init --recursive opencv && \
+    OPENCV_PYTHON_SKIP_GIT_COMMANDS=1 \
+    /opt/python/cp312-cp312/bin/pip wheel . --no-deps -w /tmp/wheelhouse && \
+    auditwheel repair /tmp/wheelhouse/opencv_python_headless-*.whl -w /wheels && \
+    rm -rf /src/opencv-python /tmp/wheelhouse
+
+# The final ffmpeg-vlm-wheel image contains (one arch per platform build):
+#   - av-*-cp311-abi3-manylinux_2_28_*.whl  (PyAV stable ABI; works on cp312)
+#   - opencv_python_headless-*-cp312-*.whl
+FROM scratch AS ffmpeg-vlm-wheel
+COPY --from=av-wheel-builder-py312 /wheels /wheels
+COPY --from=opencv-headless-wheel-builder-py312 /wheels /wheels