195 lines
7.7 KiB
Bash
Executable File
195 lines
7.7 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
# install_server.sh — one-shot server install for ComfyUI-MiMoASR.
|
|
#
|
|
# Run this on the ComfyUI server (the machine with the GPU), AFTER
|
|
# you've cloned this node pack into ComfyUI/custom_nodes/.
|
|
#
|
|
# What it does:
|
|
# 1. Clones the upstream MiMo-V2.5-ASR repo
|
|
# 2. Downloads the two model checkpoints
|
|
# 3. Pins transformers to 4.49.0 in your ComfyUI venv
|
|
# 4. Installs MiMo's other runtime deps
|
|
# 5. Builds flash-attn from source (only if not already importable)
|
|
#
|
|
# What it does NOT do:
|
|
# - Launch ComfyUI (that's your call — `python main.py --listen 0.0.0.0`)
|
|
# - Modify your shell rc files
|
|
# - Touch other custom nodes
|
|
#
|
|
# Configure via env vars before running:
|
|
# COMFYUI_DIR — where your ComfyUI lives (must have .venv/)
|
|
# MODELS_DIR — where to put MiMo weights (default: $COMFYUI_DIR/models/mimo)
|
|
# REPOS_DIR — where to clone upstream repo (default: $COMFYUI_DIR/repos)
|
|
# HF_HOME — Hugging Face cache (default: $MODELS_DIR/_hf_cache)
|
|
# CUDA_HOME — required for flash-attn build (default: /usr/local/cuda)
|
|
# GPU_ARCH — TORCH_CUDA_ARCH_LIST (default: auto-detected via torch)
|
|
# MAX_JOBS — parallel compile jobs for flash-attn (default: 4)
|
|
# SKIP_DOWNLOAD — set to 1 to skip the model download step
|
|
# SKIP_FLASH_ATTN — set to 1 to skip the flash-attn build step
|
|
|
|
set -euo pipefail
|
|
|
|
# ---------- Configuration ----------------------------------------------------
|
|
COMFYUI_DIR="${COMFYUI_DIR:-/mnt/modelall/ComfyUI}"
|
|
MODELS_DIR="${MODELS_DIR:-$COMFYUI_DIR/models/mimo}"
|
|
REPOS_DIR="${REPOS_DIR:-$COMFYUI_DIR/repos}"
|
|
HF_HOME="${HF_HOME:-$MODELS_DIR/_hf_cache}"
|
|
CUDA_HOME="${CUDA_HOME:-/usr/local/cuda}"
|
|
MAX_JOBS="${MAX_JOBS:-4}"
|
|
|
|
export HF_HOME
|
|
|
|
# ---------- Helpers ----------------------------------------------------------
|
|
log() { printf '\033[1;36m▸\033[0m %s\n' "$*"; }
|
|
ok() { printf '\033[1;32m✓\033[0m %s\n' "$*"; }
|
|
warn() { printf '\033[1;33m!\033[0m %s\n' "$*"; }
|
|
fatal() { printf '\033[1;31m✗\033[0m %s\n' "$*" >&2; exit 1; }
|
|
|
|
# ---------- Pre-flight -------------------------------------------------------
|
|
log "Configuration:"
|
|
echo " COMFYUI_DIR = $COMFYUI_DIR"
|
|
echo " MODELS_DIR = $MODELS_DIR"
|
|
echo " REPOS_DIR = $REPOS_DIR"
|
|
echo " HF_HOME = $HF_HOME"
|
|
echo " CUDA_HOME = $CUDA_HOME"
|
|
echo " MAX_JOBS = $MAX_JOBS"
|
|
echo
|
|
|
|
[ -d "$COMFYUI_DIR" ] || fatal "ComfyUI dir not found: $COMFYUI_DIR (set COMFYUI_DIR)"
|
|
[ -d "$COMFYUI_DIR/.venv" ] || fatal "No .venv in $COMFYUI_DIR — this script expects a uv-managed venv"
|
|
|
|
mkdir -p "$MODELS_DIR" "$REPOS_DIR" "$HF_HOME"
|
|
|
|
# Activate venv
|
|
# shellcheck disable=SC1091
|
|
source "$COMFYUI_DIR/.venv/bin/activate"
|
|
ok "Activated venv: $(which python)"
|
|
|
|
# ---------- 1. Clone upstream repo -------------------------------------------
|
|
log "Step 1/5: Clone XiaomiMiMo/MiMo-V2.5-ASR"
|
|
if [ -d "$REPOS_DIR/MiMo-V2.5-ASR/.git" ]; then
|
|
ok "Already cloned at $REPOS_DIR/MiMo-V2.5-ASR"
|
|
else
|
|
git clone https://github.com/XiaomiMiMo/MiMo-V2.5-ASR.git "$REPOS_DIR/MiMo-V2.5-ASR"
|
|
fi
|
|
|
|
if [ ! -d "$REPOS_DIR/MiMo-V2.5-ASR/src/mimo_audio" ]; then
|
|
warn "src/mimo_audio/ missing — running git submodule init"
|
|
(cd "$REPOS_DIR/MiMo-V2.5-ASR" && git submodule update --init --recursive)
|
|
fi
|
|
[ -d "$REPOS_DIR/MiMo-V2.5-ASR/src/mimo_audio" ] \
|
|
|| fatal "Still no src/mimo_audio/ — repo structure may have changed upstream"
|
|
ok "Upstream repo OK"
|
|
|
|
# ---------- 2. Download models -----------------------------------------------
|
|
log "Step 2/5: Download model weights"
|
|
if [ "${SKIP_DOWNLOAD:-0}" = "1" ]; then
|
|
warn "SKIP_DOWNLOAD=1, skipping"
|
|
else
|
|
uv pip install -U huggingface-hub --quiet
|
|
if [ ! -f "$MODELS_DIR/MiMo-V2.5-ASR/config.json" ]; then
|
|
log " Downloading MiMo-V2.5-ASR (~16-32 GB)"
|
|
hf download XiaomiMiMo/MiMo-V2.5-ASR --local-dir "$MODELS_DIR/MiMo-V2.5-ASR"
|
|
else
|
|
ok "MiMo-V2.5-ASR already present"
|
|
fi
|
|
if [ ! -f "$MODELS_DIR/MiMo-Audio-Tokenizer/config.json" ]; then
|
|
log " Downloading MiMo-Audio-Tokenizer"
|
|
hf download XiaomiMiMo/MiMo-Audio-Tokenizer --local-dir "$MODELS_DIR/MiMo-Audio-Tokenizer"
|
|
else
|
|
ok "MiMo-Audio-Tokenizer already present"
|
|
fi
|
|
fi
|
|
|
|
# ---------- 3. Pin transformers ----------------------------------------------
|
|
log "Step 3/5: Snapshot freeze + pin transformers==4.49.0"
|
|
FREEZE_FILE="$HOME/pip_freeze_before_mimo_$(date +%Y%m%d_%H%M%S).txt"
|
|
uv pip freeze > "$FREEZE_FILE"
|
|
ok "Snapshot saved to $FREEZE_FILE (use this to roll back if something breaks)"
|
|
|
|
current_transformers=$(python -c "import transformers; print(transformers.__version__)" 2>/dev/null || echo "none")
|
|
if [ "$current_transformers" = "4.49.0" ]; then
|
|
ok "transformers already at 4.49.0"
|
|
else
|
|
warn "Current transformers: $current_transformers — downgrading to 4.49.0"
|
|
uv pip install "transformers==4.49.0" "tokenizers>=0.21,<0.22" --force-reinstall
|
|
fi
|
|
|
|
# ---------- 4. Other deps ----------------------------------------------------
|
|
log "Step 4/5: Install MiMo runtime deps + audio I/O"
|
|
uv pip install "accelerate>=1.9.0" "librosa>=0.11.0" "scipy>=1.16.1" "zhon==2.1.1"
|
|
uv pip install torchaudio requests websocket-client
|
|
ok "Deps installed"
|
|
|
|
# ---------- 5. flash-attn ----------------------------------------------------
|
|
log "Step 5/5: flash-attn"
|
|
if [ "${SKIP_FLASH_ATTN:-0}" = "1" ]; then
|
|
warn "SKIP_FLASH_ATTN=1, skipping"
|
|
elif python -c "from flash_attn import flash_attn_varlen_func; assert flash_attn_varlen_func is not None" 2>/dev/null; then
|
|
ok "flash-attn already installed and working"
|
|
else
|
|
warn "flash-attn not importable — building from source (20-40 min)"
|
|
[ -x "$CUDA_HOME/bin/nvcc" ] || fatal "nvcc not found at $CUDA_HOME/bin/nvcc — set CUDA_HOME"
|
|
|
|
# Auto-detect GPU arch if not provided
|
|
if [ -z "${GPU_ARCH:-}" ]; then
|
|
GPU_ARCH=$(python -c "
|
|
import torch
|
|
major, minor = torch.cuda.get_device_capability(0)
|
|
print(f'{major}.{minor}')" 2>/dev/null) || GPU_ARCH="8.6"
|
|
log " Auto-detected GPU arch: $GPU_ARCH (override with GPU_ARCH=...)"
|
|
fi
|
|
|
|
uv pip install ninja packaging
|
|
|
|
export CUDA_HOME
|
|
export PATH="$CUDA_HOME/bin:$PATH"
|
|
export LD_LIBRARY_PATH="$CUDA_HOME/lib64:${LD_LIBRARY_PATH:-}"
|
|
export TORCH_CUDA_ARCH_LIST="$GPU_ARCH"
|
|
export MAX_JOBS
|
|
|
|
BUILD_LOG="$HOME/flash_attn_build_$(date +%Y%m%d_%H%M%S).log"
|
|
log " Building (this WILL take a long time). Log: $BUILD_LOG"
|
|
log " Tip: detach from screen/tmux now and run \`tail -f $BUILD_LOG\` in another shell"
|
|
|
|
if uv pip install flash-attn --no-build-isolation 2>&1 | tee "$BUILD_LOG"; then
|
|
if python -c "from flash_attn import flash_attn_varlen_func; assert flash_attn_varlen_func is not None"; then
|
|
ok "flash-attn built and importable"
|
|
else
|
|
fatal "flash-attn install reported success but import returned None. Check $BUILD_LOG"
|
|
fi
|
|
else
|
|
fatal "flash-attn build failed. See $BUILD_LOG"
|
|
fi
|
|
fi
|
|
|
|
# ---------- Done -------------------------------------------------------------
|
|
cat <<EOF
|
|
|
|
$(ok "Install complete.")
|
|
|
|
Next steps:
|
|
|
|
1. Launch ComfyUI with the env vars set:
|
|
|
|
export HF_HOME=$HF_HOME
|
|
export MIMO_ASR_REPO=$REPOS_DIR/MiMo-V2.5-ASR
|
|
cd $COMFYUI_DIR && source .venv/bin/activate
|
|
python main.py --listen 0.0.0.0 --port 8188
|
|
|
|
2. In the MiMoASRLoader node, set:
|
|
|
|
model_path = $MODELS_DIR/MiMo-V2.5-ASR
|
|
tokenizer_path = $MODELS_DIR/MiMo-Audio-Tokenizer
|
|
repo_path = (blank — uses MIMO_ASR_REPO env var)
|
|
|
|
3. From the client, smoke-test the CLI:
|
|
|
|
uv run python cli/mimo_asr_cli.py speech.wav \\
|
|
--ip SERVER_IP \\
|
|
--model-path $MODELS_DIR/MiMo-V2.5-ASR \\
|
|
--tokenizer-path $MODELS_DIR/MiMo-Audio-Tokenizer \\
|
|
--repo-path $REPOS_DIR/MiMo-V2.5-ASR
|
|
|
|
EOF
|