Dockerfile

# syntax=docker/dockerfile:1
# AudioMuse-AI Dockerfile
# Supports both CPU (ubuntu:24.04) and GPU (nvidia/cuda:13.3.1-cudnn-runtime-ubuntu24.04) builds
#
# Build examples:
#   CPU:  docker build -t audiomuse-ai .
#   GPU:  docker build --build-arg BASE_IMAGE=nvidia/cuda:13.3.1-cudnn-runtime-ubuntu24.04 -t audiomuse-ai-gpu .

ARG BASE_IMAGE=ubuntu:24.04
# Optional: URL of a prebuilt aarch64 onnxruntime-gpu wheel. There is no official
# linux/aarch64 wheel for onnxruntime-gpu on PyPI or pypi.nvidia.com, so ARM64
# (DGX Spark / GB10, compute capability sm_121) GPU builds install this URL
# instead of the onnxruntime-gpu pin in gpu.txt. Both GPU requirement files now
# target CUDA 13 / cuDNN 9: gpu.txt pins onnxruntime-gpu 1.28.0 (CUDA 13) plus
# cupy-cuda13x/cuml-cu13, and gpu-arm64.txt uses the same cupy/cuml pins with
# this prebuilt wheel. When this arg is set, BASE_IMAGE must be a CUDA 13.x image.
# Example: https://github.com/NeptuneHub/AudioMuse-AI/releases/download/v5.0.0-model/onnxruntime_gpu-1.25.0-cp312-cp312-linux_aarch64.whl
ARG ONNXRUNTIME_WHEEL_URL=""

# ============================================================================
# Stage 1: Download ML models (cached separately for faster rebuilds)
# ============================================================================
FROM ubuntu:24.04 AS models

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

RUN mkdir -p /app/model

# Install download tools with exponential backoff retry
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/*

# Download musicnn ONNX models with diagnostics and retry logic.
# Lyrics / CLAP / HuggingFace bundles are downloaded in the RUN steps below,
# all within this models stage so they are cached independently of the
# Python requirements and application code.
RUN set -eux; \
    mkdir -p /app/model; \
    urls=( \
        "https://github.com/NeptuneHub/AudioMuse-AI/releases/download/v5.0.0-model/musicnn_embedding.onnx" \
        "https://github.com/NeptuneHub/AudioMuse-AI/releases/download/v5.0.0-model/musicnn_prediction.onnx" \
        "https://github.com/NeptuneHub/AudioMuse-AI/releases/download/v5.0.0-model/neural_fingerprint.onnx" \
        "https://github.com/NeptuneHub/AudioMuse-AI/releases/download/v5.0.0-model/neural_fingerprint_pq.npz" \
    ); \
    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

# Download the HuggingFace cache tarball from the GitHub release, then trim it.
# Only the roberta-base *tokenizer* is used at runtime (the CLAP text encoder runs
# as ONNX); bert/bart and the roberta weights are stripped below to shrink the image.
RUN set -eux; \
    base_url="https://github.com/NeptuneHub/AudioMuse-AI/releases/download/v5.0.0-model"; \
    hf_models="huggingface_models.tar.gz"; \
    cache_dir="/app/.cache/huggingface"; \
    echo "Downloading HuggingFace models (~985MB)..."; \
    \
    # Download with retry logic \
    n=0; \
    until [ "$n" -ge 5 ]; do \
        if wget --no-verbose --tries=3 --retry-connrefused --waitretry=10 \
            --header="User-Agent: AudioMuse-Docker/1.0 (+https://github.com/NeptuneHub/AudioMuse-AI)" \
            -O "/tmp/$hf_models" "$base_url/$hf_models"; then \
            echo "✓ HuggingFace models downloaded"; \
            break; \
        fi; \
        n=$((n+1)); \
        echo "Download attempt $n failed - retrying in $((n*n))s"; \
        sleep $((n*n)); \
    done; \
    if [ "$n" -ge 5 ]; then \
        echo "ERROR: Failed to download HuggingFace models after 5 attempts"; \
        exit 1; \
    fi; \
    \
    # Extract to cache directory \
    mkdir -p "$cache_dir"; \
    echo "Extracting HuggingFace models..."; \
    tar -xzf "/tmp/$hf_models" -C "$cache_dir"; \
    \
    # Verify extraction \
    if [ ! -d "$cache_dir/hub" ]; then \
        echo "ERROR: HuggingFace models extraction failed"; \
        exit 1; \
    fi; \
    \
    # Trim the HF cache to just the roberta-base tokenizer (~1.4 GB saved). \
    # The app's only runtime HF dependency is AutoTokenizer.from_pretrained("roberta-base") \
    # (tasks/clap_analyzer.py); the CLAP text encoder runs as ONNX (clap_text_model.onnx), \
    # so bert-base-uncased, bart-base and the roberta model weights are never loaded. \
    # A tokenizer needs only tokenizer.json/vocab.json/merges.txt/config (all < 2 MB). \
    # NOTE: Dockerfile-noavx2 is intentionally left unchanged. \
    hub_dir="$cache_dir/hub"; \
    rm -rf "$hub_dir/models--bert-base-uncased" "$hub_dir/models--facebook--bart-base"; \
    rb="$hub_dir/models--roberta-base"; \
    if [ -d "$rb" ]; then \
        find "$rb/blobs" -type f -size +10M -delete; \
        find "$rb/snapshots" \( -name "model.safetensors" -o -name "pytorch_model.bin" \) -delete; \
    fi; \
    \
    # Clean up tarball \
    rm -f "/tmp/$hf_models"; \
    \
    echo "✓ HuggingFace models extracted and trimmed to $cache_dir"; \
    du -sh "$cache_dir"

# Download CLAP ONNX models
# - DCLAP audio model (~20MB + external data): Distilled student for music analysis in worker containers
# - Text model (~478MB): Original LAION CLAP text encoder for text search in Flask containers
RUN set -eux; \
    dclap_url="https://github.com/NeptuneHub/AudioMuse-AI-DCLAP/releases/download/v1"; \
    text_url="https://github.com/NeptuneHub/AudioMuse-AI/releases/download/v5.0.0-model"; \
    arch=$(uname -m); \
    echo "Architecture detected: $arch - Downloading CLAP ONNX models..."; \
    \
    # Download DCLAP audio model (~1.2MB ONNX + ~20MB external data) \
    n=0; \
    until [ "$n" -ge 5 ]; do \
        if wget --no-verbose --tries=3 --retry-connrefused --waitretry=10 \
            --header="User-Agent: AudioMuse-Docker/1.0 (+https://github.com/NeptuneHub/AudioMuse-AI)" \
            -O "/app/model/model_epoch_36.onnx" "$dclap_url/model_epoch_36.onnx"; then \
            echo "✓ DCLAP audio model downloaded"; \
            break; \
        fi; \
        n=$((n+1)); \
        echo "Download attempt $n for DCLAP audio model failed - retrying in $((n*n))s"; \
        sleep $((n*n)); \
    done; \
    if [ "$n" -ge 5 ]; then \
        echo "ERROR: Failed to download DCLAP audio model after 5 attempts"; \
        exit 1; \
    fi; \
    \
    # Download DCLAP audio model external data file \
    n=0; \
    until [ "$n" -ge 5 ]; do \
        if wget --no-verbose --tries=3 --retry-connrefused --waitretry=10 \
            --header="User-Agent: AudioMuse-Docker/1.0 (+https://github.com/NeptuneHub/AudioMuse-AI)" \
            -O "/app/model/model_epoch_36.onnx.data" "$dclap_url/model_epoch_36.onnx.data"; then \
            echo "✓ DCLAP audio model data downloaded"; \
            break; \
        fi; \
        n=$((n+1)); \
        echo "Download attempt $n for DCLAP audio data failed - retrying in $((n*n))s"; \
        sleep $((n*n)); \
    done; \
    if [ "$n" -ge 5 ]; then \
        echo "ERROR: Failed to download DCLAP audio model data after 5 attempts"; \
        exit 1; \
    fi; \
    \
    # Download text model (~478MB) \
    text_model="clap_text_model.onnx"; \
    n=0; \
    until [ "$n" -ge 5 ]; do \
        if wget --no-verbose --tries=3 --retry-connrefused --waitretry=10 \
            --header="User-Agent: AudioMuse-Docker/1.0 (+https://github.com/NeptuneHub/AudioMuse-AI)" \
            -O "/app/model/$text_model" "$text_url/$text_model"; then \
            echo "✓ CLAP text model downloaded"; \
            break; \
        fi; \
        n=$((n+1)); \
        echo "Download attempt $n for text model failed - retrying in $((n*n))s"; \
        sleep $((n*n)); \
    done; \
    if [ "$n" -ge 5 ]; then \
        echo "ERROR: Failed to download CLAP text model after 5 attempts"; \
        exit 1; \
    fi; \
    \
    # Verify DCLAP audio model \
    if [ ! -f "/app/model/model_epoch_36.onnx" ]; then \
        echo "ERROR: DCLAP audio model file not created"; \
        exit 1; \
    fi; \
    if [ ! -f "/app/model/model_epoch_36.onnx.data" ]; then \
        echo "ERROR: DCLAP audio model data file not created"; \
        exit 1; \
    fi; \
    \
    # Verify text model \
    if [ ! -f "/app/model/$text_model" ]; then \
        echo "ERROR: CLAP text model file not created"; \
        exit 1; \
    fi; \
    file_size=$(stat -c%s "/app/model/$text_model" 2>/dev/null || stat -f%z "/app/model/$text_model" 2>/dev/null || echo "0"); \
    if [ "$file_size" -lt 450000000 ]; then \
        echo "ERROR: CLAP text model file is too small (expected ~478MB, got $file_size bytes)"; \
        exit 1; \
    fi; \
    \
    echo "✓ CLAP models downloaded successfully (arch: $arch)"; \
    ls -lh /app/model/model_epoch_36.onnx /app/model/model_epoch_36.onnx.data "/app/model/$text_model"

# Download the SAE concept-steering graphs (~2 MB each)
# The encoder maps a DCLAP embedding to 1024 concept latents, the decoder maps
# them back. The concept catalogue that names those latents is small and ships
# inside the repository, so only the two graphs are fetched here.
RUN set -eux; \
    sae_url="https://github.com/NeptuneHub/AudioMuse-AI-SAE/releases/download/v1"; \
    for sae_model in dclap_sae_k20_d1024_best_encoder.onnx dclap_sae_k20_d1024_best_decoder.onnx; do \
        n=0; \
        until [ "$n" -ge 5 ]; do \
            if wget --no-verbose --tries=3 --retry-connrefused --waitretry=10 \
                --header="User-Agent: AudioMuse-Docker/1.0 (+https://github.com/NeptuneHub/AudioMuse-AI)" \
                -O "/app/model/$sae_model" "$sae_url/$sae_model"; then \
                echo "SAE model $sae_model downloaded"; \
                break; \
            fi; \
            n=$((n+1)); \
            echo "Download attempt $n for $sae_model failed - retrying in $((n*n))s"; \
            sleep $((n*n)); \
        done; \
        if [ "$n" -ge 5 ]; then \
            echo "ERROR: Failed to download $sae_model after 5 attempts"; \
            exit 1; \
        fi; \
        if [ ! -f "/app/model/$sae_model" ]; then \
            echo "ERROR: $sae_model not created"; \
            exit 1; \
        fi; \
    done; \
    echo "SAE models downloaded successfully"; \
    ls -lh /app/model/dclap_sae_k20_d1024_best_encoder.onnx /app/model/dclap_sae_k20_d1024_best_decoder.onnx

# Download Whisper-small ONNX bundle (~570 MB) - HuggingFace optimum export
# of openai/whisper-small (encoder_model.onnx + decoder_model_merged.onnx +
# tokenizer files + preprocessor config). Re-hosted on the project's GitHub
# release for mirror independence. Bundle ships `whisper-small-onnx/` as
# its top-level directory. Loaded at runtime by lyrics/whisper_onnx.py via
# raw onnxruntime.
RUN set -eux; \
    whisper_dir="/app/model/whisper-small-onnx"; \
    whisper_url="https://github.com/NeptuneHub/AudioMuse-AI/releases/download/v5.0.0-model/lyrics_model_whisper.tar.gz"; \
    whisper_dest="/tmp/lyrics_model_whisper.tar.gz"; \
    echo "Downloading Whisper-small ONNX bundle (~570 MB)..."; \
    n=0; \
    until [ "$n" -ge 5 ]; do \
        if wget --no-verbose --tries=3 --retry-connrefused --waitretry=10 \
            --header="User-Agent: AudioMuse-Docker/1.0 (+https://github.com/NeptuneHub/AudioMuse-AI)" \
            -O "$whisper_dest" "$whisper_url"; then \
            echo "✓ whisper bundle downloaded"; break; \
        fi; \
        n=$((n+1)); \
        echo "wget attempt $n for whisper bundle failed - retrying in $((n*n))s"; \
        sleep $((n*n)); \
    done; \
    if [ "$n" -ge 5 ]; then \
        echo "ERROR: failed to download whisper bundle"; exit 1; \
    fi; \
    mkdir -p /app/model; \
    tar -xzf "$whisper_dest" -C /app/model; \
    rm -f "$whisper_dest"; \
    for f in encoder_model.onnx decoder_model_merged.onnx \
             tokenizer.json tokenizer_config.json \
             special_tokens_map.json preprocessor_config.json \
             config.json generation_config.json vocab.json merges.txt; do \
        if [ ! -f "$whisper_dir/$f" ]; then \
            echo "ERROR: Whisper file missing: $whisper_dir/$f"; \
            echo "Actual /app/model contents:"; \
            ls -laR /app/model | head -50; \
            exit 1; \
        fi; \
    done; \
    echo "✓ Whisper-small ONNX model ready in $whisper_dir"; \
    du -sh "$whisper_dir"

# Download silero VAD ONNX (~2 MB) - re-hosted on the project's GitHub release
# for mirror independence (original source: snakers4/silero-vad). Bundle ships
# silero_vad.onnx at archive root. Loaded by lyrics/silero_onnx.py via raw
# onnxruntime.
RUN set -eux; \
    silero_url="https://github.com/NeptuneHub/AudioMuse-AI/releases/download/v5.0.0-model/lyrics_model_silero_vad.tar.gz"; \
    silero_dest="/tmp/lyrics_model_silero_vad.tar.gz"; \
    silero_path="/app/model/silero_vad.onnx"; \
    echo "Downloading silero VAD ONNX bundle (~2 MB)..."; \
    n=0; \
    until [ "$n" -ge 5 ]; do \
        if wget --no-verbose --tries=3 --retry-connrefused --waitretry=5 \
            --header="User-Agent: AudioMuse-Docker/1.0 (+https://github.com/NeptuneHub/AudioMuse-AI)" \
            -O "$silero_dest" "$silero_url"; then \
            echo "✓ silero bundle downloaded"; break; \
        fi; \
        n=$((n+1)); \
        echo "wget attempt $n for silero bundle failed - retrying in $((n*n))s"; \
        sleep $((n*n)); \
    done; \
    if [ "$n" -ge 5 ]; then \
        echo "ERROR: failed to download silero bundle"; exit 1; \
    fi; \
    mkdir -p /app/model; \
    tar -xzf "$silero_dest" -C /app/model; \
    rm -f "$silero_dest"; \
    if [ ! -f "$silero_path" ]; then \
        echo "ERROR: silero_vad.onnx missing after extraction"; \
        ls -laR /app/model | head -50; \
        exit 1; \
    fi; \
    ls -lh "$silero_path"

# Download gte-multilingual-base INT8 ONNX bundle (~325 MB) - multilingual
# sentence embedding pre-exported and dynamic-INT8-quantized by this project.
# Tarball ships the ONNX file flat at the archive root
# (`gte-multilingual-base-int8.onnx`) plus a sibling `gte-multilingual-base/`
# directory with the tokenizer files. Loaded by lyrics/gte_onnx.py via raw
# onnxruntime + the bare `tokenizers` package (CLS pooling + L2 norm at runtime).
RUN set -eux; \
    gte_onnx_path="/app/model/gte-multilingual-base-int8.onnx"; \
    gte_tok_dir="/app/model/gte-multilingual-base"; \
    gte_url="https://github.com/NeptuneHub/AudioMuse-AI/releases/download/v5.0.0-model/lyrics_model_gte_vnni.tar.gz"; \
    gte_dest="/tmp/lyrics_model_gte_vnni.tar.gz"; \
    echo "Downloading gte-multilingual-base INT8 ONNX bundle (~325 MB)..."; \
    n=0; \
    until [ "$n" -ge 5 ]; do \
        if wget --no-verbose --tries=3 --retry-connrefused --waitretry=10 \
            --header="User-Agent: AudioMuse-Docker/1.0 (+https://github.com/NeptuneHub/AudioMuse-AI)" \
            -O "$gte_dest" "$gte_url"; then \
            echo "✓ gte bundle downloaded"; break; \
        fi; \
        n=$((n+1)); \
        echo "wget attempt $n for gte bundle failed - retrying in $((n*n))s"; \
        sleep $((n*n)); \
    done; \
    if [ "$n" -ge 5 ]; then \
        echo "ERROR: failed to download gte bundle"; exit 1; \
    fi; \
    mkdir -p /app/model; \
    tar -xzf "$gte_dest" -C /app/model; \
    rm -f "$gte_dest"; \
    if [ ! -f "$gte_onnx_path" ]; then \
        echo "ERROR: gte ONNX missing after extraction: $gte_onnx_path"; \
        ls -laR /app/model | head -50; \
        exit 1; \
    fi; \
    for f in tokenizer.json tokenizer_config.json config.json special_tokens_map.json; do \
        if [ ! -f "$gte_tok_dir/$f" ]; then \
            echo "ERROR: gte tokenizer file missing: $gte_tok_dir/$f"; exit 1; \
        fi; \
    done; \
    echo "✓ gte-multilingual-base ONNX ready ($gte_onnx_path + $gte_tok_dir)"; \
    du -sh "$gte_onnx_path" "$gte_tok_dir"

# ============================================================================
# Stage 2a: runtime-base - RUNTIME-ONLY system libs (parent of `runner`)
# ============================================================================
# This stage holds only what the application needs at run time: shared
# libraries (.so) that Python wheels load, plus the small set of CLI tools
# the entrypoint / supervisord / debugging rely on. It deliberately omits
# compilers and -dev headers - those live in the `base` stage below, which
# is used solely to build Python wheels in the `libraries` stage and never
# becomes a parent of `runner`.
#
# `cuda-compiler` is INTENTIONALLY kept here (not moved to build-only)
# because cupy JIT-compiles CUDA kernels at runtime on GPU builds.
FROM ${BASE_IMAGE} AS runtime-base

ARG BASE_IMAGE

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

RUN set -ux; \
    n=0; \
    until [ "$n" -ge 5 ]; do \
        # Use noninteractive frontend to avoid tzdata prompts when installing tzdata
        if DEBIAN_FRONTEND=noninteractive apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
            python3 python3-pip \
            libfftw3-double3=3.3.10-1ubuntu3 \
            libyaml-0-2=0.2.5-1build1 \
            libsamplerate0=0.2.2-4build1 \
            libsndfile1=1.2.2-1ubuntu5.24.04.1 \
            libopenblas0 \
            liblapack3=3.12.0-3build1.1 \
            libgomp1 \
            libpq5 \
            ffmpeg libchromaprint-tools wget curl \
            supervisor procps \
            git vim strace iputils-ping \
            postgresql-common ca-certificates \
            "$(if [[ "$BASE_IMAGE" =~ ^nvidia/cuda:([0-9]+)\.([0-9]+).+$ ]]; then echo "cuda-compiler-${BASH_REMATCH[1]}-${BASH_REMATCH[2]}"; fi)" \
            # PostgreSQL 18 client from PGDG (pg_dump 18 backs up PG 15-18; psql restore stays compatible with old pg_dump 16 / PG 15 dumps)
            && /usr/share/postgresql-common/pgdg/apt.postgresql.org.sh -y \
            && DEBIAN_FRONTEND=noninteractive apt-get update \
            && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends postgresql-client-18; 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 && \
    rm -f /usr/lib/python3.*/EXTERNALLY-MANAGED

# ============================================================================
# Stage 2b: base - runtime-base + compilers / -dev headers (BUILD-ONLY)
# ============================================================================
# Adds the toolchain needed to compile Python wheels (psycopg2, essentia,
# numpy/scipy fallbacks, etc.). Parent of `libraries` only - `runner`
# branches off `runtime-base`, so gcc/g++/python3-dev and the -dev headers
# never reach the final published image.
FROM runtime-base AS base

ARG BASE_IMAGE

# Copy uv for fast package management (10-100x faster than pip)
COPY --from=ghcr.io/astral-sh/uv:latest /uv /bin/uv

RUN set -ux; \
    n=0; \
    until [ "$n" -ge 5 ]; do \
        if DEBIAN_FRONTEND=noninteractive apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
            python3-dev \
            libfftw3-dev \
            libyaml-dev \
            libsamplerate0-dev \
            libsndfile1-dev \
            libopenblas-dev \
            liblapack-dev=3.12.0-3build1.1 \
            libpq-dev \
            gcc g++; 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/*

# ============================================================================
# Stage 3: Libraries - Python packages installation
# ============================================================================
FROM base AS libraries

ARG BASE_IMAGE
ARG ONNXRUNTIME_WHEEL_URL

WORKDIR /app

# Copy requirements files
COPY requirements/ /app/requirements/

# Install Python packages with uv (combined in single layer for efficiency)
# GPU builds: cupy, cuml, onnxruntime-gpu, torch (CUDA)
# CPU builds: onnxruntime (CPU only), torch (CPU)
# Note: --index-strategy unsafe-best-match resolves conflicts between pypi.nvidia.com and pypi.org
RUN rm -f /usr/lib/python3.*/EXTERNALLY-MANAGED; \
    export UV_BREAK_SYSTEM_PACKAGES=1; \
    if [[ "$BASE_IMAGE" =~ ^nvidia/cuda: ]]; then \
        if [[ -n "$ONNXRUNTIME_WHEEL_URL" ]]; then \
            echo "NVIDIA (arm64/CUDA 13) base: installing GPU packages (cupy-cuda13x, cuml-cu13) with prebuilt aarch64 onnxruntime-gpu wheel"; \
            uv pip install --system --no-cache --index-strategy unsafe-best-match -r /app/requirements/gpu-arm64.txt -r /app/requirements/common.txt "$ONNXRUNTIME_WHEEL_URL" || exit 1; \
        else \
            echo "NVIDIA base image detected: installing GPU packages (cupy, cuml, onnxruntime-gpu, torch+cuda)"; \
            uv pip install --system --no-cache --index-strategy unsafe-best-match -r /app/requirements/gpu.txt -r /app/requirements/common.txt || exit 1; \
        fi \
    else \
        echo "CPU base image: installing all packages together for dependency resolution"; \
        uv pip install --system --no-cache --index-strategy unsafe-best-match -r /app/requirements/cpu.txt -r /app/requirements/common.txt || exit 1; \
    fi \
    && echo "Verifying psycopg2 installation..." \
    && python3 -c "import psycopg2; print('psycopg2 OK')" \
    && find /usr/local/lib/python3.*/dist-packages -type d -name __pycache__ -exec rm -rf {} + 2>/dev/null || true \
    && find /usr/local/lib/python3.*/dist-packages -type f \( -name "*.pyc" -o -name "*.pyo" \) -delete \
    && PYV=$(python3 -c 'import sys; print("python%d.%d" % sys.version_info[:2])') \
    && mkdir -p /app/site-packages && cp -a /usr/local/lib/$PYV/dist-packages/. /app/site-packages/ \
    && du -sh /app/site-packages

# ============================================================================
# Stage 4: Runner - Final production image
# ============================================================================
# IMPORTANT: extends `runtime-base` (NOT `base`). That keeps gcc/g++,
# python3-dev and the *-dev headers out of the final image, saving
# ~300-400 MB. Anything that needs compiling lives in the `libraries`
# stage and gets COPY'd in as already-built artifacts below.
FROM runtime-base AS runner

ENV LANG=C.UTF-8 \
    PYTHONUNBUFFERED=1 \
    DEBIAN_FRONTEND=noninteractive \
    TZ=UTC \
    IVF_DISK_CACHE_DIR=/app/ivf_cache \
    HF_HOME=/app/.cache/huggingface \
    HF_HUB_DISABLE_XET=1 \
    HF_XET_DISABLE=1 \
    GUNICORN_CMD_ARGS="--config /app/gunicorn.conf.py"

# Note: bundled HuggingFace models (RoBERTa, ...) load with
# local_files_only=True per call. The gte/whisper/silero ONNX bundles are
# pre-downloaded as release tarballs; HF_HUB_OFFLINE is intentionally NOT set.

WORKDIR /app

# Ensure tzdata package is installed so /usr/share/zoneinfo exists and TZ can be applied
RUN set -eux; \
    apt-get update && apt-get install -y --no-install-recommends tzdata && rm -rf /var/lib/apt/lists/*

# Copy all downloaded/extracted models from the models stage
COPY --from=models /app/model/ /app/model/
# Copy HuggingFace cache (text encoders) from the models stage
COPY --from=models /app/.cache/huggingface/ /app/.cache/huggingface/

# Verify cache was copied correctly
RUN ls -lah /app/.cache/huggingface/ && \
    echo "HuggingFace cache contents:" && \
    du -sh /app/.cache/huggingface/* || echo "Cache directory empty!"

# Copy Python packages from libraries stage
COPY --from=libraries /app/site-packages/ /tmp/site-packages/
RUN PYV=$(python3 -c 'import sys; print("python%d.%d" % sys.version_info[:2])') && \
    mkdir -p /usr/local/lib/$PYV/dist-packages && \
    cp -a /tmp/site-packages/. /usr/local/lib/$PYV/dist-packages/ && \
    rm -rf /tmp/site-packages
# Copy console entrypoints (gunicorn, etc.) from libraries stage
COPY --from=libraries /usr/local/bin/ /usr/local/bin/

# Copy application code (last to maximize cache hits for code changes)
COPY . /app
COPY deployment/docker-entrypoint.sh /app/docker-entrypoint.sh
COPY deployment/supervisord.conf /etc/supervisor/conf.d/supervisord.conf
RUN chmod +x /app/docker-entrypoint.sh
RUN ls -l /etc/supervisor/conf.d && test -f /etc/supervisor/conf.d/supervisord.conf

# ============================================================================
# CPU CONSISTENCY SETTINGS
# ============================================================================
# These environment variables ensure CONSISTENT behavior across different
# AVX2-capable CPUs (e.g., Intel 6th gen vs 12th gen have different FPU defaults).
# They do NOT enable non-AVX support - AVX2 is still required for x86_64 builds.
# ARM64 builds use NEON instructions and work on all ARM64 CPUs.

# oneDNN floating-point math mode: STRICT reduces non-deterministic FP optimizations
# Keeps CPU behavior deterministic across different CPU generations
ENV ONEDNN_DEFAULT_FPMATH_MODE=STRICT

# ONNX Runtime optimization settings to prevent signal 9 crashes on newer CPUs
# (Intel 12600K and similar have different optimization behavior than older CPUs)
# Similar to TF_ENABLE_ONEDNN_OPTS=0 for TensorFlow compatibility
ENV ORT_DISABLE_ALL_OPTIMIZATIONS=1 \
    ORT_ENABLE_CPU_FP16_OPS=0

# Force consistent memory allocation and precision behavior
# Prevents different memory allocation patterns and floating-point precision issues
# between Intel generations (e.g., 12600K vs i5-6500)
ENV ORT_DISABLE_AVX512=1 \
    ORT_FORCE_SHARED_PROVIDER=1

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

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

# numba JIT cache must land in a writable directory.
# When the container runs as a non-root user the system site-packages directory
# (/usr/local/lib/python3.x/dist-packages/) is read-only, which causes librosa
# to fail with: "cannot cache function: no locator available".
# Point numba to /tmp so it always has write access (issue: NeptuneHub/AudioMuse-AI#479).
ENV NUMBA_CACHE_DIR=/tmp/numba_cache

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

EXPOSE 8000

WORKDIR /workspace
ENTRYPOINT ["/app/docker-entrypoint.sh"]
CMD []