summaryrefslogtreecommitdiff
path: root/makima/sh/download-models.sh
blob: 4f188f31abebc989ce16c5c93ef86c4863349a15 (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
#!/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 "$@"