summaryrefslogtreecommitdiff
path: root/makima/sh/download-models.sh
blob: 4f1bd77f488fb2c2f90799badd0ca6bb0e1365c0 (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
#!/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 Chatterbox TTS models (for TTS functionality)
CHATTERBOX_MODEL_DIR="${CHATTERBOX_MODEL_DIR:-/app/models/chatterbox-turbo}"

download_chatterbox_tts() {
    if [ -d "$CHATTERBOX_MODEL_DIR" ] && \
       [ -f "$CHATTERBOX_MODEL_DIR/speech_encoder.onnx" ] && \
       [ -f "$CHATTERBOX_MODEL_DIR/embed_tokens.onnx" ] && \
       [ -f "$CHATTERBOX_MODEL_DIR/language_model.onnx" ] && \
       [ -f "$CHATTERBOX_MODEL_DIR/conditional_decoder.onnx" ] && \
       [ -f "$CHATTERBOX_MODEL_DIR/tokenizer.json" ]; then
        echo "Chatterbox TTS models already exist, skipping..."
        return 0
    fi

    echo "Downloading Chatterbox TTS models..."
    mkdir -p "$CHATTERBOX_MODEL_DIR"

    # Download only the required files from ResembleAI/chatterbox-turbo-ONNX
    # (skip fp16, q4, q4f16, quantized variants to save ~4GB)
    echo "Downloading ResembleAI/chatterbox-turbo-ONNX (standard precision only)..."
    local tmpdir=$(mktemp -d)
    hf download ResembleAI/chatterbox-turbo-ONNX \
        --include "onnx/speech_encoder.onnx" \
        --include "onnx/speech_encoder.onnx_data" \
        --include "onnx/embed_tokens.onnx" \
        --include "onnx/embed_tokens.onnx_data" \
        --include "onnx/language_model.onnx" \
        --include "onnx/language_model.onnx_data" \
        --include "onnx/conditional_decoder.onnx" \
        --include "onnx/conditional_decoder.onnx_data" \
        --include "tokenizer.json" \
        --local-dir "$tmpdir"

    # Move files to the model directory (flatten the onnx/ subdirectory)
    mv "$tmpdir"/onnx/*.onnx "$CHATTERBOX_MODEL_DIR"/
    mv "$tmpdir"/onnx/*.onnx_data "$CHATTERBOX_MODEL_DIR"/
    mv "$tmpdir"/tokenizer.json "$CHATTERBOX_MODEL_DIR"/
    rm -rf "$tmpdir"

    echo "Chatterbox TTS models downloaded successfully"
}

download_chatterbox_tts

# Execute the main command
exec "$@"