#!/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 <