[FEEDBACK] How to use gpu?

[pmffromspace]

Hey awesome project - I got relativly large collection and wanted to use gpu to speed scan speed up:

In docker-compose.yml I tried added nvidia container toolkit support to it like so but it doesnt seem to automatically switch to gpu if available. Do I need to set a env var? couldnt find any infos :)

#worker
...
    runtime: nvidia
    restart: unless-stopped
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities:
                - gpu

[NeptuneHub]

Thanks for the feedback. I'm not an expert of nvidia gpu, because in my homelab I scale up on the CPU (I have a cluster with multiple node but no CPU), anyway a contributors times ago created this deployment/docker-compose-worker-nvidia.yaml that also use the specific containers ghcr.io/neptunehub/audiomuse-ai:latest-nvidia.

Have a look at it, and if you spot any missing configuration feel free to raise a PR!

[pmffromspace]

Hey thanks for that fast reply!

Unfortunatly even with this docker compose and the tag latest-nvidia it does not use the gpu. The container can access the gpu no problem (I tested the "nvidia-smi" command inside the container and it shows the gpu"

I will investigate further...

(Just so Im not on the wrong path: the scanning should already use gpu or only the clustering?)

[NeptuneHub]

For scanning I use ONNX, it should use gpu if present but I don’t know if some env variable is needed. I never investigated because I don’t have a gpu even for test.

About clustering I don’t know if GPU is supposed to be used. I use scikit-learn for clustering.

GPU worked in the old Tensorflow implementation (v0.6.*-beta) probably some change are needed for ONNX ?

[aimbot-luc]

I also noticed that only the cpu is being used instead of my nvidia gpu. Tried 0.7.4-beta-nvidia with the additions of docker-compose-worker-nvidia.yaml. Only this if I'm not blind:

NVIDIA_VISIBLE_DEVICES: "0"
NVIDIA_DRIVER_CAPABILITIES: "compute,utility"

...

    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]

I tried to replace onnxruntime to onnxruntime-gpu inside the container, no luck yet.. But I don't know what I'm doing though haha ;)

Do you know a way to check whether the gpu is initialize/used inside the container?

[NeptuneHub]

Looking at the documentation here:
https://onnxruntime.ai/docs/install/#cuda-and-cudnn

you're right need the onnxruntime-gpu, so this mean that build-nvidia.yml and Dockrfile need to be adapated for this. If you find a way to have it work please raise a PR.

To look if the GPU is used inside the container I think you first of all look if it used out of the container. If is not used outside, then you need to throubleshoot the container itself. But unluky I'm not an expert and I don't have actually a machine to test&learn.

With Tensorflow @ravenscroftj and @HeroBrine1st helped me in creating the -nvidia image, maybe they are so kind to help us also with ONNX.

[HeroBrine1st]

I can currently leave a hint that Dockerfile already uses conditional on CUDA:

      "$(if [[ "$BASE_IMAGE" =~ ^nvidia/cuda:([0-9]+)\.([0-9]+).+$ ]]; then echo "cuda-compiler-${BASH_REMATCH[1]}-${BASH_REMATCH[2]}"; fi)" \

It can be used in pip install too, just needs else/fi for CPU image and capture groups (with their contents) can be removed (so after : there's only .+$). Don't forget about SHELL set to /bin/bash

That'll be one-line change (UPD: oops, plus SHELL instruction 😆)

[NeptuneHub]

Checking the onx documentation and going of cut&paste in the Dockerfile I producted something like the code below, where I just put specific dependencies and onnx gpu runtime in the if.
The point is that I would like to avoid to push a change that I'm not able to test.

# or nvidia/cuda:12.8.1-cudnn-runtime-ubuntu22.04 (CUDA version can go as low as CUDA 12.2 but need to check)
ARG BASE_IMAGE=ubuntu:22.04

FROM ubuntu:22.04 AS models

SHELL ["/bin/bash", "-lc"]

RUN mkdir -p /app/model

RUN set -ux; \
  n=0; \
  until [ "$n" -ge 5 ]; do \
    if apt-get update && apt-get install -y --no-install-recommends wget ca-certificates curl; then \
      break; \
    fi; \
    n=$((n+1)); \
    echo "apt-get attempt $n failed — retrying in $((n*n))s"; \
    sleep $((n*n)); \
  done; \
  rm -rf /var/lib/apt/lists/*

RUN set -eux; \
  urls=( \
    "https://github.com/NeptuneHub/AudioMuse-AI/releases/download/v2.0.0-model/danceability-msd-musicnn-1.onnx" \
    "https://github.com/NeptuneHub/AudioMuse-AI/releases/download/v2.0.0-model/mood_aggressive-msd-musicnn-1.onnx" \
    "https://github.com/NeptuneHub/AudioMuse-AI/releases/download/v2.0.0-model/mood_happy-msd-musicnn-1.onnx" \
    "https://github.com/NeptuneHub/AudioMuse-AI/releases/download/v2.0.0-model/mood_party-msd-musicnn-1.onnx" \
    "https://github.com/NeptuneHub/AudioMuse-AI/releases/download/v2.0.0-model/mood_relaxed-msd-musicnn-1.onnx" \
    "https://github.com/NeptuneHub/AudioMuse-AI/releases/download/v2.0.0-model/mood_sad-msd-musicnn-1.onnx" \
    "https://github.com/NeptuneHub/AudioMuse-AI/releases/download/v2.0.0-model/msd-msd-musicnn-1.onnx" \
    "https://github.com/NeptuneHub/AudioMuse-AI/releases/download/v2.0.0-model/msd-musicnn-1.onnx" \
  ); \
  mkdir -p /app/model; \
  for u in "${urls[@]}"; do \
    n=0; \
    fname="/app/model/$(basename "$u")"; \
    # Diagnostic: print server response headers (helpful when downloads return 0 bytes)
    wget --server-response --spider --timeout=15 --header="User-Agent: AudioMuse-Docker/1.0 (+https://github.com/NeptuneHub/AudioMuse-AI)" "$u" || true; \
    until [ "$n" -ge 5 ]; do \
      # Use wget with retries. --tries and --waitretry add backoff for transient failures.
      if wget --no-verbose --tries=3 --retry-connrefused --waitretry=5 --header="User-Agent: AudioMuse-Docker/1.0 (+https://github.com/NeptuneHub/AudioMuse-AI)" -O "$fname" "$u"; then \
        echo "Downloaded $u -> $fname"; \
        break; \
      fi; \
      n=$((n+1)); \
      echo "wget attempt $n for $u failed — retrying in $((n*n))s"; \
      sleep $((n*n)); \
    done; \
    if [ "$n" -ge 5 ]; then \
      echo "ERROR: failed to download $u after 5 attempts"; \
      ls -lah /app/model || true; \
      exit 1; \
    fi; \
  done

FROM ${BASE_IMAGE} AS base

ARG BASE_IMAGE

SHELL ["/bin/bash", "-c"]

# Install system dependencies, including ffmpeg which is crucial for pydub
# cuda-compiler is only for libdevice.10.bc, can be extracted into another stage
RUN set -ux; \
  n=0; \
  until [ "$n" -ge 5 ]; do \
    if apt-get update && apt-get install -y --no-install-recommends \
      python3 python3-pip python3-dev \
      libfftw3-3=3.3.8-2ubuntu8 libyaml-0-2 libsamplerate0 \
      libsndfile1=1.0.31-2ubuntu0.2 \
      ffmpeg wget git vim \
      redis-tools curl \
      supervisor \
      strace \
      procps \
      iputils-ping \
      libopenblas-dev=0.3.20+ds-1 \
      liblapack-dev=3.10.0-2ubuntu1 \
      libpq-dev \
      gcc \
      g++ \
      "$(if [[ "$BASE_IMAGE" =~ ^nvidia/cuda:([0-9]+)\.([0-9]+).+$ ]]; then echo "cuda-compiler-${BASH_REMATCH[1]}-${BASH_REMATCH[2]}"; fi)"; then \
      break; \
    fi; \
    n=$((n+1)); \
    echo "apt-get attempt $n failed — retrying in $((n*n))s"; \
    sleep $((n*n)); \
  done; \
  rm -rf /var/lib/apt/lists/*

#RUN test -f /usr/local/cuda-12.8/nvvm/libdevice/libdevice.10.bc

FROM base AS libraries

ARG TARGETARCH
ARG BASE_IMAGE

# pydub is for audio conversion
# Pin numpy to a stable version to avoid numeric differences between builds
RUN --mount=type=cache,target=/root/.cache/pip \
    bash -lc '\
      # Conditionally install GPU-only helper packages if using an NVIDIA base image\
      GPU_PKGS="flatbuffers packaging protobuf sympy"; \
      # Then install the main package set (without onnxruntime and without GPU-only pkgs)\
      pip3 install --prefix=/install --no-cache-dir \
        numpy==1.26.4 \
        scipy==1.15.3 \
        numba==0.60.0 \
        soundfile==0.13.1 \
        Flask \
        Flask-Cors \
        redis \
        requests \
        scikit-learn==1.7.2 \
        rq \
        pyyaml \
        six \
        voyager==2.1.0 \
        rapidfuzz \
        psycopg2-binary \
        ftfy \
        flasgger \
        sqlglot \
        google-generativeai \
        mistralai \
        umap-learn \
        pydub \
        python-mpd2 \
        onnx==1.14.1 \
        librosa==0.11.0 || exit 1; \
      # Conditionally install the ONNX Runtime wheel and GPU-only helper packages
      if [[ "${BASE_IMAGE}" =~ ^nvidia/cuda: ]]; then \
        echo "Detected NVIDIA base image: installing GPU-only packages and onnxruntime-gpu from custom index"; \
        pip3 install --prefix=/install --no-cache-dir $GPU_PKGS || exit 1; \
        pip3 install --prefix=/install --no-cache-dir --index-url "https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-11/pypi/simple/" onnxruntime-gpu==1.15.1 || exit 1; \
      else \
        echo "CPU base image: installing onnxruntime (CPU) only"; \
        pip3 install --prefix=/install --no-cache-dir onnxruntime==1.15.1 || exit 1; \
      fi'

FROM base AS runner

ENV LANG=C.UTF-8 \
    PYTHONUNBUFFERED=1 \
    DEBIAN_FRONTEND=noninteractive

WORKDIR /app

COPY --from=libraries /install/ /usr/
COPY --from=models /app/model/ /app/model/

COPY . /app

COPY deployment/supervisord.conf /etc/supervisor/conf.d/supervisord.conf

# oneDNN floating-point math mode: STRICT reduces non-deterministic FP optimizations
# Keep this if you want more deterministic CPU behavior when using oneDNN-enabled runtimes
ENV ONEDNN_DEFAULT_FPMATH_MODE=STRICT

# ONNX Runtime optimization settings to prevent signal 9 crashes on newer CPUs (Intel 12600K, etc.)
# Similar to TF_ENABLE_ONEDNN_OPTS=0 for TensorFlow compatibility
ENV ORT_DISABLE_ALL_OPTIMIZATIONS=1
ENV ORT_ENABLE_CPU_FP16_OPS=0

# Force consistent memory allocation and precision behavior (Intel 12600K vs i5-6500 compatibility)
# Prevents different memory allocation patterns and floating-point precision issues
ENV ORT_DISABLE_AVX512=1
ENV ORT_FORCE_SHARED_PROVIDER=1

# Force consistent floating-point behavior across different Intel generations
# 12600K has different FPU precision defaults than 6th gen CPUs
ENV MKL_ENABLE_INSTRUCTIONS=AVX2
ENV MKL_DYNAMIC=FALSE

# Prevent aggressive memory pre-allocation on newer CPUs
ENV ORT_DISABLE_MEMORY_PATTERN_OPTIMIZATION=1

ENV PYTHONPATH=/usr/local/lib/python3/dist-packages:/app

EXPOSE 8000

WORKDIR /workspace
CMD ["bash", "-c", "if [ \"$SERVICE_TYPE\" = \"worker\" ]; then echo 'Starting worker processes via supervisord...' && /usr/bin/supervisord -c /etc/supervisor/conf.d/supervisord.conf; else echo 'Starting web service...' && python3 /app/app.py; fi"]

[pmffromspace]

So I build the image using this Dockerfile (no error in build) but it still does not seem to use the gpu when doing analysis

[aimbot-luc]

I built it yesterday from the provided Dockerfile as well. At first, I also didn't get the GPU to work. After some error fighting I got it working, but it was not as fast as I was hoping for (around 2x faster than CPU). I did the following (probably forgot some things):

• Install nvidia-container-toolkit on host
• Image build command using the CUDA Ubuntu as base image: DOCKER_BUILDKIT=1 docker build -t audiomuse-ai-flask-app:gpu --build-arg BASE_IMAGE=nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04 .
• This link didn't work: https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/onnxruntime-cuda-11/pypi/simple/. Removed --index-url "htts://... as onnxruntime-gpu==1.15.1 is available on PyPI. See Dockerfile.
• Log is spamming Numba needs NumPy 2.0 or less. Got NumPy 2.2. when starting an analysis job. I think NumPY 2.2 is installed by default in the CUDA base image. First remove NumPY 2.2 by adding this to Dockerfile:

  apt-get remove -y python3-numpy || true; \
  apt-get autoremove -y || true;

• Random installation location issues: removed --prefix=/install from pip3 install. Added:
  COPY --from=libraries /usr/local/lib/python3.10/dist-packages/ /usr/local/lib/python3.10/dist-packages/ instead of COPY --from=libraries /install/ /usr/ in Dockerfile.
• ONNX Runtime had trouble finding its CUDA provider, fixed it by adding ['CUDAExecutionProvider', 'CPUExecutionProvider'] at some points in analysis.py. See analysis.py lines 346:406 snippet below Dockerfile.
• At this point GPU started to work.
• I got more errors after, but thats on me. 1080ti VRAM is completely used by more than 2 workers... I started at least 10 for my first test run...
• Also added my docker compose below the .py file.
• Hopefully this gives some clues, or the most hacky and shit solution... :)
  |
  |
  |
  |
• Updated Dockerfile:

# or nvidia/cuda:12.8.1-cudnn-runtime-ubuntu22.04 (CUDA version can go as low as CUDA 12.2 but need to check)
ARG BASE_IMAGE=ubuntu:22.04
# ARG BASE_IMAGE=nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04

FROM ubuntu:22.04 AS models

SHELL ["/bin/bash", "-lc"]

RUN mkdir -p /app/model

RUN set -ux; \
  n=0; \
  until [ "$n" -ge 5 ]; do \
    if apt-get update && apt-get install -y --no-install-recommends wget ca-certificates curl; then \
      break; \
    fi; \
    n=$((n+1)); \
    echo "apt-get attempt $n failed — retrying in $((n*n))s"; \
    sleep $((n*n)); \
  done; \
  rm -rf /var/lib/apt/lists/*

RUN set -eux; \
  urls=( \
    "https://github.com/NeptuneHub/AudioMuse-AI/releases/download/v2.0.0-model/danceability-msd-musicnn-1.onnx" \
    "https://github.com/NeptuneHub/AudioMuse-AI/releases/download/v2.0.0-model/mood_aggressive-msd-musicnn-1.onnx" \
    "https://github.com/NeptuneHub/AudioMuse-AI/releases/download/v2.0.0-model/mood_happy-msd-musicnn-1.onnx" \
    "https://github.com/NeptuneHub/AudioMuse-AI/releases/download/v2.0.0-model/mood_party-msd-musicnn-1.onnx" \
    "https://github.com/NeptuneHub/AudioMuse-AI/releases/download/v2.0.0-model/mood_relaxed-msd-musicnn-1.onnx" \
    "https://github.com/NeptuneHub/AudioMuse-AI/releases/download/v2.0.0-model/mood_sad-msd-musicnn-1.onnx" \
    "https://github.com/NeptuneHub/AudioMuse-AI/releases/download/v2.0.0-model/msd-msd-musicnn-1.onnx" \
    "https://github.com/NeptuneHub/AudioMuse-AI/releases/download/v2.0.0-model/msd-musicnn-1.onnx" \
  ); \
  mkdir -p /app/model; \
  for u in "${urls[@]}"; do \
    n=0; \
    fname="/app/model/$(basename "$u")"; \
    # Diagnostic: print server response headers (helpful when downloads return 0 bytes)
    wget --server-response --spider --timeout=15 --header="User-Agent: AudioMuse-Docker/1.0 (+https://github.com/NeptuneHub/AudioMuse-AI)" "$u" || true; \
    until [ "$n" -ge 5 ]; do \
      # Use wget with retries. --tries and --waitretry add backoff for transient failures.
      if wget --no-verbose --tries=3 --retry-connrefused --waitretry=5 --header="User-Agent: AudioMuse-Docker/1.0 (+https://github.com/NeptuneHub/AudioMuse-AI)" -O "$fname" "$u"; then \
        echo "Downloaded $u -> $fname"; \
        break; \
      fi; \
      n=$((n+1)); \
      echo "wget attempt $n for $u failed — retrying in $((n*n))s"; \
      sleep $((n*n)); \
    done; \
    if [ "$n" -ge 5 ]; then \
      echo "ERROR: failed to download $u after 5 attempts"; \
      ls -lah /app/model || true; \
      exit 1; \
    fi; \
  done

FROM ${BASE_IMAGE} AS base

ARG BASE_IMAGE

SHELL ["/bin/bash", "-c"]

# Install system dependencies, including ffmpeg which is crucial for pydub
# cuda-compiler is only for libdevice.10.bc, can be extracted into another stage
RUN set -ux; \
  n=0; \
  until [ "$n" -ge 5 ]; do \
    if apt-get update && apt-get install -y --no-install-recommends \
      python3 python3-pip python3-dev \
      libfftw3-3=3.3.8-2ubuntu8 libyaml-0-2 libsamplerate0 \
      libsndfile1=1.0.31-2ubuntu0.2 \
      ffmpeg wget git vim \
      redis-tools curl \
      supervisor \
      strace \
      procps \
      iputils-ping \
      libopenblas-dev=0.3.20+ds-1 \
      liblapack-dev=3.10.0-2ubuntu1 \
      libpq-dev \
      gcc \
      g++ \
      "$(if [[ "$BASE_IMAGE" =~ ^nvidia/cuda:([0-9]+)\.([0-9]+).+$ ]]; then echo "cuda-compiler-${BASH_REMATCH[1]}-${BASH_REMATCH[2]}"; fi)"; then \
      break; \
    fi; \
    n=$((n+1)); \
    echo "apt-get attempt $n failed — retrying in $((n*n))s"; \
    sleep $((n*n)); \
  done; \
  rm -rf /var/lib/apt/lists/* \
  apt-get remove -y python3-numpy || true; \
  apt-get autoremove -y || true;

#RUN test -f /usr/local/cuda-12.8/nvvm/libdevice/libdevice.10.bc

FROM base AS libraries

ARG TARGETARCH
ARG BASE_IMAGE

# pydub is for audio conversion
# Pin numpy to a stable version to avoid numeric differences between builds
RUN --mount=type=cache,target=/root/.cache/pip \
    bash -lc '\
      GPU_PKGS="flatbuffers packaging protobuf sympy"; \
      pip3 install --no-cache-dir numpy==1.26.4 || exit 1; \
      pip3 install --no-cache-dir \
        scipy==1.15.3 \
        numba==0.60.0 \
        soundfile==0.13.1 \
        Flask \
        Flask-Cors \
        redis \
        requests \
        scikit-learn==1.7.2 \
        rq \
        pyyaml \
        six \
        voyager==2.1.0 \
        rapidfuzz \
        psycopg2-binary \
        ftfy \
        flasgger \
        sqlglot \
        google-generativeai \
        mistralai \
        umap-learn \
        pydub \
        python-mpd2 \
        onnx==1.14.1 \
        librosa==0.11.0 || exit 1; \
      if [[ "${BASE_IMAGE}" =~ ^nvidia/cuda: ]]; then \
        echo "Detected NVIDIA base image: installing GPU-only packages and onnxruntime-gpu"; \
        pip3 install --no-cache-dir $GPU_PKGS || exit 1; \
        pip3 install --no-cache-dir onnxruntime-gpu==1.15.1 || exit 1; \
      else \
        echo "CPU base image: installing onnxruntime (CPU) only"; \
        pip3 install --no-cache-dir onnxruntime==1.15.1 || exit 1; \
      fi'


FROM base AS runner

ENV LANG=C.UTF-8 \
    PYTHONUNBUFFERED=1 \
    DEBIAN_FRONTEND=noninteractive

WORKDIR /app

# COPY --from=libraries /install/ /usr/
COPY --from=libraries /usr/local/lib/python3.10/dist-packages/ /usr/local/lib/python3.10/dist-packages/

COPY --from=models /app/model/ /app/model/

COPY . /app

COPY deployment/supervisord.conf /etc/supervisor/conf.d/supervisord.conf

# oneDNN floating-point math mode: STRICT reduces non-deterministic FP optimizations
# Keep this if you want more deterministic CPU behavior when using oneDNN-enabled runtimes
ENV ONEDNN_DEFAULT_FPMATH_MODE=STRICT

# ONNX Runtime optimization settings to prevent signal 9 crashes on newer CPUs (Intel 12600K, etc.)
# Similar to TF_ENABLE_ONEDNN_OPTS=0 for TensorFlow compatibility
ENV ORT_DISABLE_ALL_OPTIMIZATIONS=1
ENV ORT_ENABLE_CPU_FP16_OPS=0

# Force consistent memory allocation and precision behavior (Intel 12600K vs i5-6500 compatibility)
# Prevents different memory allocation patterns and floating-point precision issues
ENV ORT_DISABLE_AVX512=1
ENV ORT_FORCE_SHARED_PROVIDER=1

# Force consistent floating-point behavior across different Intel generations
# 12600K has different FPU precision defaults than 6th gen CPUs
ENV MKL_ENABLE_INSTRUCTIONS=AVX2
ENV MKL_DYNAMIC=FALSE

# Prevent aggressive memory pre-allocation on newer CPUs
ENV ORT_DISABLE_MEMORY_PATTERN_OPTIMIZATION=1

ENV PYTHONPATH=/usr/local/lib/python3/dist-packages:/app

EXPOSE 8000

WORKDIR /workspace
CMD ["bash", "-c", "if [ \"$SERVICE_TYPE\" = \"worker\" ]; then echo 'Starting worker processes via supervisord...' && /usr/bin/supervisord -c /etc/supervisor/conf.d/supervisord.conf; else echo 'Starting web service...' && python3 /app/app.py; fi"]

|
|
|
|

• Updated analysis.py at line 350, 360 and 385:

# --- 3. Run Main Models (Embedding and Prediction) ---
    try:
        # Load and run embedding model (ONNX)
        # embedding_sess = ort.InferenceSession(model_paths['embedding'])
        # ORIGINAL!!!: embedding_sess = ort.InferenceSession(model_paths['embedding'])
        embedding_sess = ort.InferenceSession(
            model_paths['embedding'],
            providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
        )
        embedding_feed_dict = {DEFINED_TENSOR_NAMES['embedding']['input']: final_patches}
        embeddings_per_patch = run_inference(embedding_sess, embedding_feed_dict, DEFINED_TENSOR_NAMES['embedding']['output'])

        # Load and run prediction model (ONNX)
        # prediction_sess = ort.InferenceSession(model_paths['prediction'])
        # ALSO CHANGED!!!
        prediction_sess = ort.InferenceSession(
            model_paths['prediction'],
            providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
        )
        prediction_feed_dict = {DEFINED_TENSOR_NAMES['prediction']['input']: embeddings_per_patch}
        mood_logits = run_inference(prediction_sess, prediction_feed_dict, DEFINED_TENSOR_NAMES['prediction']['output'])

        averaged_logits = np.mean(mood_logits, axis=0)
        # Apply sigmoid to convert raw model outputs (logits) into probabilities
        final_mood_predictions = sigmoid(averaged_logits)

        moods = {label: float(score) for label, score in zip(mood_labels_list, final_mood_predictions)}

    except Exception as e:
        logger.error(f"Main model inference failed for {os.path.basename(file_path)}: {e}", exc_info=True)
        return None, None
        
    # --- 4. Run Secondary Models ---
    other_predictions = {}

    for key in ["danceable", "aggressive", "happy", "party", "relaxed", "sad"]:
        try:
            model_path = model_paths[key]
            # other_sess = ort.InferenceSession(model_path)
            # ALSO CHANGED !!!
            other_sess = ort.InferenceSession(
                model_path,
                providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
            )
            feed_dict = {DEFINED_TENSOR_NAMES[key]['input']: embeddings_per_patch}
            probabilities_per_patch = run_inference(other_sess, feed_dict, DEFINED_TENSOR_NAMES[key]['output'])

            if probabilities_per_patch is None:
                other_predictions[key] = 0.0
            else:
                if isinstance(probabilities_per_patch, np.ndarray) and probabilities_per_patch.ndim == 2 and probabilities_per_patch.shape[1] == 2:
                    # Using the CLASS_INDEX_MAP to select the correct probability
                    positive_class_index = CLASS_INDEX_MAP.get(key, 0)
                    class_probs = probabilities_per_patch[:, positive_class_index]
                    other_predictions[key] = float(np.mean(class_probs))
                else:
                    other_predictions[key] = 0.0

        except Exception as e:
            logger.error(f"Error predicting '{key}' for {os.path.basename(file_path)}: {e}", exc_info=True)
            other_predictions[key] = 0.0

|
|
|
|

• Docker compose file:

services:
################
## AUDIO MUSE
################
# docker compose up -d --scale audiomuse-ai-worker=2
  redis:
    image: redis:7-alpine
    container_name: audiomuse-redis
    ports:
      - "6379:6379" # Expose Redis port to the host
    volumes:
      - redis-data:/data # Persistent storage for Redis data
    restart: unless-stopped

  # PostgreSQL database service
  postgres:
    image: postgres:15-alpine
    container_name: audiomuse-postgres
    environment:
      POSTGRES_USER: "audiomuse"
      POSTGRES_PASSWORD: "audiomusepassword"
      POSTGRES_DB: "audiomusedb"
    ports:
      - "5432:5432" # Expose PostgreSQL port to the host
    volumes:
      - postgres-data:/var/lib/postgresql/data # Persistent storage for PostgreSQL data
    restart: unless-stopped

  # AudioMuse-AI Flask application service
  audiomuse-ai-flask:
    image: audiomuse-ai-flask-app:gpu
#    image: ghcr.io/neptunehub/audiomuse-ai:0.7.4-beta-nvidia
#    runtime: nvidia
    container_name: audiomuse-ai-flask-app
    ports:
      - "8654:8000"
    environment:
      SERVICE_TYPE: "flask"
      MEDIASERVER_TYPE: "jellyfin"
      JELLYFIN_URL: "http://xxxx"
      JELLYFIN_USER_ID: "xxxx"
      JELLYFIN_TOKEN: "xxxx"
      POSTGRES_USER: "audiomuse"
      POSTGRES_PASSWORD: "audiomusepassword"
      POSTGRES_DB: "audiomusedb"
      POSTGRES_HOST: "postgres"
      POSTGRES_PORT: "5432"
      REDIS_URL: "redis://redis:6379/0"
      GEMINI_API_KEY: "YOUR_GEMINI_API_KEY_HERE"
      TEMP_DIR: "/app/temp_audio"
    volumes:
      - temp-audio-flask:/app/temp_audio
    depends_on:
      - redis
      - postgres
    restart: unless-stopped
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]

  # AudioMuse-AI RQ Worker service
  audiomuse-ai-worker:
    image: audiomuse-ai-flask-app:gpu
#    runtime: nvidia
    environment:
      SERVICE_TYPE: "worker"
      MEDIASERVER_TYPE: "jellyfin"
      JELLYFIN_URL: "http://xxxx"
      JELLYFIN_USER_ID: "xxxxx"
      JELLYFIN_TOKEN: "xxxxx"
      POSTGRES_USER: "audiomuse"
      POSTGRES_PASSWORD: "audiomusepassword"
      POSTGRES_DB: "audiomusedb"
      POSTGRES_HOST: "postgres"
      POSTGRES_PORT: "5432"
      REDIS_URL: "redis://redis:6379/0"
      GEMINI_API_KEY: "YOUR_GEMINI_API_KEY_HERE"
      TEMP_DIR: "/app/temp_audio"
    volumes:
      - temp-audio-worker:/app/temp_audio
    depends_on:
      - redis
      - postgres
    restart: unless-stopped
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: all
              capabilities: [gpu]


# Define volumes for persistent data and temporary files
volumes:
  redis-data:
  postgres-data:
  temp-audio-flask: # Volume for Flask app's temporary audio
  temp-audio-worker: # Volume for Worker's temporary audio

[NeptuneHub]

Can you raise directly a PR with the Dockerfile, analysis.py updated and also the update of deployment/docker-compose-nvidia.yaml if you did update also on it ?

So that I can do some regression test also to look that on CPU image there is no issue and then merge.
Thanks !

[pmffromspace]

I can replicate this setup and confirm that gpu is now in use. But looking at my logs I can see some errors and also all analysis tasks seem to not complete.

@aimbot-luc do you get similar errors or is this a problem on my end?

(important to use docker build -t gpu-test --build-arg BASE_IMAGE=nvidia/cuda:11.8.0-cudnn8-runtime-ubuntu22.04 .)

audiomuse-ai-worker-1   | [CRITICAL]-[28-10-2025 08-01-39]-Album analysis REDACTED failed: cannot unpack non-iterable NoneType object
audiomuse-ai-worker-1   | Traceback (most recent call last):
audiomuse-ai-worker-1   |   File "/app/tasks/analysis.py", line 499, in analyze_album_task
audiomuse-ai-worker-1   |     analysis, embedding = analyze_track(path, MOOD_LABELS, model_paths)
audiomuse-ai-worker-1   | TypeError: cannot unpack non-iterable NoneType object
audiomuse-ai-worker-1   | [INFO]-[28-10-2025 08-01-39]-[AlbumTask-cbc81a37-e4fe-40f3-a501-bc7af7a85e38-Tom Misch - Geography (2018) [24B-44.1kHz]] Failed to analyze album 'REDACTED': cannot unpack non-iterable NoneType object
audiomuse-ai-worker-1   | [ERROR]-[28-10-2025 08-01-39]-Worker ee7e6a1a74b1433fbbcd5f223dea02ad: job cbc81a37-e4fe-40f3-a501-bc7af7a85e38: exception raised while executing (tasks.analysis.analyze_album_task)
audiomuse-ai-worker-1   | Traceback (most recent call last):
audiomuse-ai-worker-1   |   File "/usr/local/lib/python3.10/dist-packages/rq/worker.py", line 1659, in perform_job
audiomuse-ai-worker-1   |     return_value = job.perform()
audiomuse-ai-worker-1   |   File "/usr/local/lib/python3.10/dist-packages/rq/job.py", line 1318, in perform
audiomuse-ai-worker-1   |     self._result = self._execute()
audiomuse-ai-worker-1   |   File "/usr/local/lib/python3.10/dist-packages/rq/job.py", line 1376, in _execute
audiomuse-ai-worker-1   |     result = self.func(*self.args, **self.kwargs)
audiomuse-ai-worker-1   |   File "/app/tasks/analysis.py", line 499, in analyze_album_task
audiomuse-ai-worker-1   |     analysis, embedding = analyze_track(path, MOOD_LABELS, model_paths)
audiomuse-ai-worker-1   | TypeError: cannot unpack non-iterable NoneType object
audiomuse-ai-worker-1   |
audiomuse-ai-worker-1   | [INFO]-[28-10-2025 08-01-39]-Job cbc81a37-e4fe-40f3-a501-bc7af7a85e38: enqueued for retry, 2 remaining

About gpu for clustering: it might also be possible but I currently dont have the time to dive into this, I will just leave this here if someone smart enough comes along ;)

Im running one worker on 4060ti (16gb vram) and it looks like this (therese is also a whisper service using some static vram..)

[aimbot-luc]

Yeah, I have the same problem, so maybe wait until someone comes along who knows what they're doing. Sorry for my late reply.