path: root/makima/sh/download-models.sh

#!/bin/bash
set -e

MODELS_DIR="${MODELS_DIR:-/app/models}"
MODEL_BASE_URL="${MODEL_BASE_URL:-}"

# Model directories
PARAKEET_DIR="parakeet-tdt-0.6b-v3"
EOU_DIR="realtime_eou_120m-v1-onnx"
DIARIZATION_DIR="diarization"

download_from_url() {
    local name=$1
    local url=$2
    local dest="$MODELS_DIR/$name"

    if [ -d "$dest" ] && [ "$(ls -A $dest 2>/dev/null)" ]; then
        echo "Model $name already exists, skipping..."
        return 0
    fi

    echo "Downloading $name from URL..."
    mkdir -p "$dest"
    curl -L "$url" | tar -xz -C "$dest" --strip-components=1
    echo "Downloaded $name successfully"
}

download_from_hf() {
    local dest=$1
    local repo=$2

    if [ -d "$dest" ] && [ "$(ls -A $dest 2>/dev/null)" ]; then
        echo "Model $dest already exists, skipping..."
        return 0
    fi

    echo "Downloading from Hugging Face ($repo)..."
    mkdir -p "$dest"
    hf download "$repo" --local-dir "$dest"
    echo "Downloaded to $dest successfully"
}

download_from_hf_subdir() {
    local dest=$1
    local repo=$2
    local subdir=$3

    if [ -d "$dest" ] && [ "$(ls -A $dest 2>/dev/null)" ]; then
        echo "Model $dest already exists, skipping..."
        return 0
    fi

    echo "Downloading $subdir from Hugging Face ($repo)..."
    local tmpdir=$(mktemp -d)
    hf download "$repo" --include "$subdir/*" --local-dir "$tmpdir"

    # Move subdirectory contents to destination
    mkdir -p "$dest"
    mv "$tmpdir/$subdir"/* "$dest"/
    rm -rf "$tmpdir"
    echo "Downloaded to $dest successfully"
}

# Check if models exist
# TDT: encoder-model.onnx, encoder-model.onnx.data, decoder_joint-model.onnx, vocab.txt
# EOU: encoder.onnx, decoder_joint.onnx, tokenizer.json
# Diarization: diar_streaming_sortformer_4spk-v2.1.onnx
check_models_exist() {
    # TDT model files
    [ -f "$MODELS_DIR/$PARAKEET_DIR/encoder-model.onnx" ] && \
    [ -f "$MODELS_DIR/$PARAKEET_DIR/decoder_joint-model.onnx" ] && \
    [ -f "$MODELS_DIR/$PARAKEET_DIR/vocab.txt" ] && \
    # EOU model files
    [ -f "$MODELS_DIR/$EOU_DIR/encoder.onnx" ] && \
    [ -f "$MODELS_DIR/$EOU_DIR/decoder_joint.onnx" ] && \
    [ -f "$MODELS_DIR/$EOU_DIR/tokenizer.json" ] && \
    # Diarization model
    [ -f "$MODELS_DIR/$DIARIZATION_DIR/diar_streaming_sortformer_4spk-v2.1.onnx" ]
}

if check_models_exist; then
    echo "All models present"
else
    mkdir -p "$MODELS_DIR"

    if [ -n "$MODEL_BASE_URL" ]; then
        echo "Downloading models from custom URL..."
        download_from_url "$PARAKEET_DIR" "${MODEL_BASE_URL}/parakeet-tdt-0.6b-v3.tar.gz"
        download_from_url "$EOU_DIR" "${MODEL_BASE_URL}/realtime_eou_120m-v1-onnx.tar.gz"
        download_from_url "$DIARIZATION_DIR" "${MODEL_BASE_URL}/diarization.tar.gz"
    else
        echo "Downloading models from Hugging Face..."

        # Parakeet TDT from istupakov/parakeet-tdt-0.6b-v3-onnx
        # Required: encoder-model.onnx, encoder-model.onnx.data, decoder_joint-model.onnx, vocab.txt
        download_from_hf "$MODELS_DIR/$PARAKEET_DIR" "istupakov/parakeet-tdt-0.6b-v3-onnx"

        # Verify TDT files were downloaded
        if [ ! -f "$MODELS_DIR/$PARAKEET_DIR/vocab.txt" ]; then
            echo "ERROR: vocab.txt not found in parakeet TDT model"
            echo "Contents of $MODELS_DIR/$PARAKEET_DIR:"
            ls -la "$MODELS_DIR/$PARAKEET_DIR"
            exit 1
        fi

        # EOU model from altunenes/parakeet-rs (subdirectory)
        download_from_hf_subdir "$MODELS_DIR/$EOU_DIR" "altunenes/parakeet-rs" "realtime_eou_120m-v1-onnx"

        # Diarization model from altunenes/parakeet-rs (single file at root)
        mkdir -p "$MODELS_DIR/$DIARIZATION_DIR"
        hf download "altunenes/parakeet-rs" "diar_streaming_sortformer_4spk-v2.1.onnx" --local-dir "$MODELS_DIR/$DIARIZATION_DIR"
    fi

    echo "All models downloaded successfully"
fi

# Download Qwen3-TTS models (for TTS functionality)
QWEN3_TTS_DIR="${QWEN3_TTS_DIR:-/app/models/qwen3-tts}"

download_qwen3_tts() {
    if [ -d "$QWEN3_TTS_DIR" ] && \
       [ -f "$QWEN3_TTS_DIR/model.safetensors" ] && \
       [ -f "$QWEN3_TTS_DIR/speech_tokenizer.safetensors" ] && \
       [ -f "$QWEN3_TTS_DIR/vocab.json" ] && \
       [ -f "$QWEN3_TTS_DIR/merges.txt" ] && \
       [ -f "$QWEN3_TTS_DIR/config.json" ]; then
        echo "Qwen3-TTS models already exist, skipping..."
        return 0
    fi

    echo "Downloading Qwen3-TTS models..."
    mkdir -p "$QWEN3_TTS_DIR"

    # Download base TTS model files from Qwen/Qwen3-TTS-12Hz-0.6B-Base
    # Note: This repo uses vocab.json + merges.txt (not tokenizer.json)
    echo "Downloading Qwen3-TTS-12Hz-0.6B-Base..."
    hf download Qwen/Qwen3-TTS-12Hz-0.6B-Base \
        model.safetensors \
        config.json \
        vocab.json \
        merges.txt \
        tokenizer_config.json \
        --local-dir "$QWEN3_TTS_DIR"

    # Download speech tokenizer from Qwen/Qwen3-TTS-Tokenizer-12Hz
    echo "Downloading Qwen3-TTS-Tokenizer-12Hz..."
    local tmpdir=$(mktemp -d)
    hf download Qwen/Qwen3-TTS-Tokenizer-12Hz \
        model.safetensors \
        --local-dir "$tmpdir"
    mv "$tmpdir/model.safetensors" "$QWEN3_TTS_DIR/speech_tokenizer.safetensors"
    rm -rf "$tmpdir"

    echo "Qwen3-TTS models downloaded successfully"
}

download_qwen3_tts

# Execute the main command
exec "$@"