commit a80d2b1910587f4b0052910320299d363f3a746e Author: hua zhuang <30758919@qq.com> Date: Thu May 21 18:48:03 2026 +0800 Initial commit: ComfyUI-MiMoASR diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..569984a --- /dev/null +++ b/.gitignore @@ -0,0 +1,38 @@ +# Python +__pycache__/ +*.py[cod] +*.so +*.egg-info/ +.eggs/ +build/ +dist/ + +# Virtualenvs +.venv/ +venv/ +env/ + +# Editor +.vscode/ +.idea/ +*.swp + +# OS +.DS_Store +Thumbs.db + +# Test / scratch +/out/ +/tmp/ +*.tar.gz +*.zip + +# Models — never commit weights +/models/ +*.safetensors +*.bin +*.pt +*.ckpt + +# CLI temp +mimoasr_*.wav diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..25c7bfa --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,71 @@ +# Changelog + +All notable changes to this project will be documented in this file. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/), +and this project adheres to [Semantic Versioning](https://semver.org/). + +## [Unreleased] + +## [0.1.2] - 2026-05-21 + +### Added +- `scripts/install_server.sh` — idempotent one-shot install for Linux + servers. Clones the upstream repo, downloads both checkpoints, pins + `transformers` to 4.49.0, installs MiMo's runtime deps, and builds + `flash-attn` from source if needed. Configurable via env vars + (`COMFYUI_DIR`, `MODELS_DIR`, `REPOS_DIR`, `GPU_ARCH`, `MAX_JOBS`). + +### Changed +- **Documentation overhaul** based on real-world install pain. + - README now front-loads the three things that will break your install + (`transformers` pin, `flash-attn` requirement, repo clone) before + any step-by-step instructions. + - Each install step gets a copy-pasteable verification command. + - Added a "compatibility matrix" gut-check table. +- **INSTALL_NOTES.md** rewritten as an exhaustive error → fix reference, + covering every error message we observed in real installs: + - `MiMo ASR repo path is empty` + - `Did not find src/mimo_audio/` + - `model path does not exist` + - `MimoAudio.__init__() got an unexpected keyword argument 'tokenizer_path'` + - `AttributeError: ... no attribute 'all_tied_weights_keys'` + - `TypeError: 'NoneType' object is not callable` (flash-attn missing) + - `flash-attn` build failures (missing nvcc, OOM, wrong arch) +- Documented the three install strategies for handling the + `transformers==4.49.0` requirement when it conflicts with other + custom nodes (accept, dedicated ComfyUI instance, roll back). +- Detailed `flash-attn` source-build guide with the exact env vars + needed for an RTX 3090 host (`TORCH_CUDA_ARCH_LIST="8.6"`, `MAX_JOBS=4`). + +### Notes +- No code changes since v0.1.1. This release is documentation-only. + +## [0.1.1] - 2026-05-21 + +### Fixed +- `MiMoASRLoader` now calls `MimoAudio(model_path, tokenizer_path)` with + positional args. The HF model card shows kwargs, but the actual + upstream `__init__` signature doesn't accept `tokenizer_path=` as a + keyword — `run_mimo_asr.py` confirms positional usage. Symptom was + `TypeError: MimoAudio.__init__() got an unexpected keyword argument 'tokenizer_path'`. + +## [0.1.0] - 2026-05-21 + +### Added +- `MiMoASRLoader` node — loads `MiMo-V2.5-ASR` + `MiMo-Audio-Tokenizer` + with process-wide caching keyed by `(model_path, tokenizer_path, repo_path)`. +- `MiMoASRTranscribe` node — runs `MimoAudio.asr_sft` on a ComfyUI + `AUDIO` input, writing it to a tempfile first; supports `Auto` / + `Chinese` / `English` language tags. +- `cli/mimo_asr_cli.py` — single-file CLI that uploads an audio file, + builds the API-format workflow, runs it over WebSocket, and prints + + optionally saves the transcript. +- Example workflow JSON at `example_workflows/mimo_asr_basic.json`. +- `INSTALL_NOTES.md` with the compatibility matrix and the + `transformers` / `flash-attn` workarounds. + +[Unreleased]: https://github.com/aadebuger/ComfyUI-MiMoASR/compare/v0.1.2...HEAD +[0.1.2]: https://github.com/aadebuger/ComfyUI-MiMoASR/compare/v0.1.1...v0.1.2 +[0.1.1]: https://github.com/aadebuger/ComfyUI-MiMoASR/compare/v0.1.0...v0.1.1 +[0.1.0]: https://github.com/aadebuger/ComfyUI-MiMoASR/releases/tag/v0.1.0 diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..c3808dc --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,41 @@ +# Contributing + +Thanks for your interest! Small contributions are very welcome. + +## Filing issues + +Useful info to include: + +- ComfyUI commit (`git -C ComfyUI rev-parse HEAD`) +- `torch.__version__`, `transformers.__version__`, GPU model +- The exact `MIMO_ASR_REPO` you're pointing at + its commit hash +- Full traceback (not just the final line) + +## Pull requests + +1. Fork + branch from `main`. +2. Keep PRs scoped — one feature or fix per PR. +3. Match the existing style (the loader / transcribe nodes show the + intended pattern: lazy heavy imports, clear runtime errors, a + `_MODEL_CACHE` for anything expensive). +4. If you add a new node, add it to **both** `__init__.py` mapping dicts + and update the README's node table. +5. Bump `CHANGELOG.md` under `[Unreleased]`. + +## Local sanity check + +```bash +# Lint +python -m py_compile nodes/*.py cli/*.py + +# Optional: install into a ComfyUI checkout and reload +ln -s "$(pwd)" /path/to/ComfyUI/custom_nodes/ComfyUI-MiMoASR +``` + +## Scope + +This pack is intentionally minimal — just the two nodes needed to use +MiMo-V2.5-ASR from ComfyUI. Anything that belongs upstream (model code, +sampler changes, etc.) should go to +[XiaomiMiMo/MiMo-V2.5-ASR](https://github.com/XiaomiMiMo/MiMo-V2.5-ASR) +instead. diff --git a/INSTALL_NOTES.md b/INSTALL_NOTES.md new file mode 100644 index 0000000..fecb11e --- /dev/null +++ b/INSTALL_NOTES.md @@ -0,0 +1,287 @@ +# Install notes + +The exhaustive companion to the README. Every error message we hit +during real-world install, with the fix and why it happens. + +If you're hitting an issue and just want it fixed, jump to +**[Error → Fix table](#error--fix-table)**. + +--- + +## Why this model is annoying to install + +Three structural issues, all in the upstream repo: + +1. **`MimoAudio` is not pip-installable.** The repo lays the code out as + `src/mimo_audio/mimo_audio.py` and uses imports of the form + `from src.mimo_audio.foo import bar`. That only resolves if the repo + *root* is on `sys.path`. The loader handles this via either the + `repo_path` node field or the `MIMO_ASR_REPO` env var. + +2. **`MimoAudio.__init__` takes positional args**, not keywords. The HF + model card shows `MimoAudio(model_path="...", tokenizer_path="...")` + but the actual signature doesn't accept those kwargs. + `run_mimo_asr.py` calls it positionally. v0.1.1+ of this node pack + uses positional args to match. + +3. **Hard dependencies on specific `transformers` and `flash-attn` + versions.** The audio tokenizer encoder calls + `flash_attn_varlen_func` with no SDPA fallback, and the modeling code + uses the pre-`all_tied_weights_keys` API from `transformers <=4.49`. + Both must be addressed before the model can run. + +--- + +## Compatibility matrix + +| Component | Upstream pin | What you probably have | What you need | +|---|---|---|---| +| Python | 3.12 | 3.10–3.13 | 3.10+ (3.13 works) | +| `torch` | 2.6.0 | 2.4–2.11 | any recent, don't downgrade | +| `transformers` | **4.49.0** | 4.50–5.x | **must downgrade to 4.49.0** | +| `tokenizers` | implicit | varies | `>=0.21,<0.22` | +| `flash-attn` | 2.7.4.post1 | none | any working build | +| CUDA toolkit | ≥12.0 | 12.4–12.6 | matches `nvcc --version` | + +--- + +## Disk space (HF cache) + +`MiMo-V2.5-ASR` is 8B params. The safetensors are stored in F32 on the +Hub (~32 GB), and HF stages a full copy in `HF_HOME` before symlinking +to your `--local-dir`. If your root partition is small: + +```bash +# Add to ~/.bashrc so every shell sees it +export HF_HOME=/big_disk/hf_cache +mkdir -p "$HF_HOME" +``` + +Symptom of not doing this: `OSError: [Errno 28] No space left on device` +half-way through the download. + +--- + +## `transformers` downgrade — what to expect + +```bash +uv pip freeze > ~/pip_freeze_before_mimo_$(date +%Y%m%d).txt +uv pip install "transformers==4.49.0" "tokenizers>=0.21,<0.22" --force-reinstall +``` + +This will almost certainly affect other ComfyUI custom nodes. + +| Custom node | Risk | Notes | +|---|---|---| +| Pure SDXL / Florence2 workflows | Low | Uses `diffusers`, mostly unaffected. | +| ComfyUI-VibeVoiceASR | Medium | Qwen2 backbone; may need adjustment. | +| ComfyUI-PaddleOCRVL | High | PaddleOCR-VL expects newer `transformers`. | +| Sonic / Qwen3TTS-class nodes | Very high | Often pinned to `transformers >=4.57`. | +| Anything with HF auto-class registration in 5.x | Very high | API renames. | + +**Strategy A: accept the conflict.** Use one ComfyUI install for MiMo +ASR, restart with the other version when you need the other nodes. + +**Strategy B: dedicated ComfyUI for MiMo.** Recommended if you use the +high-risk nodes daily. + +```bash +# Clone the venv (cheap because Python deps share .so files via inodes) +cd /mnt/modelall +cp -r ComfyUI ComfyUI-MiMo +cd ComfyUI-MiMo && source .venv/bin/activate +uv pip install "transformers==4.49.0" "tokenizers>=0.21,<0.22" --force-reinstall + +# Launch on a different port +python main.py --listen 0.0.0.0 --port 8189 +``` + +Then point the CLI at the dedicated instance: `--ip SERVER --port 8189`. + +**Strategy C: roll back if the downgrade breaks too much.** + +```bash +uv pip install -r ~/pip_freeze_before_mimo_*.txt --force-reinstall +``` + +--- + +## Building `flash-attn` from source + +The precompiled `flash-attn==2.7.4.post1` wheel only matches one exact +combo: **`cu12 + torch 2.6 + Python 3.12 + cxx11abi=False`**. Anything +else needs a source build. + +### Prerequisites + +```bash +# Tooling +uv pip install ninja packaging + +# CUDA toolkit must be present and findable. +# Default Ubuntu location is /usr/local/cuda-XX.X — adjust to yours. +ls /usr/local/cuda-12.6/bin/nvcc # must exist +``` + +If `nvcc` is missing, install the matching toolkit: + +```bash +sudo apt install cuda-toolkit-12-6 # adjust to your CUDA version +``` + +### Build + +```bash +cd ComfyUI && source .venv/bin/activate + +export CUDA_HOME=/usr/local/cuda-12.6 +export PATH=$CUDA_HOME/bin:$PATH +export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH + +# GPU compute capability — set this to avoid building for archs you don't have +# 3090 = 8.6, A100 = 8.0, 4090 = 8.9, H100 = 9.0 +export TORCH_CUDA_ARCH_LIST="8.6" + +# Parallel jobs. Each nvcc job uses ~2 GB RAM. Tune to your hardware. +export MAX_JOBS=4 + +# Build (will take 20-40 min on a 3090 host) +nohup uv pip install flash-attn --no-build-isolation \ + > ~/flash_attn_build.log 2>&1 & +tail -f ~/flash_attn_build.log +``` + +When it finishes, the log ends with something like +`Successfully installed flash-attn-X.Y.Z`. Verify: + +```bash +python -c "from flash_attn import flash_attn_varlen_func; print(flash_attn_varlen_func)" +# Must print: +``` + +### Common build failures + +| Log message | Fix | +|---|---| +| `nvcc: command not found` | `export PATH=$CUDA_HOME/bin:$PATH` | +| `cannot find -lcuda` | `export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH` | +| `g++: fatal error: cuda_runtime.h: No such file` | Install matching `cuda-toolkit-XX-X` | +| `nvcc fatal: Unsupported gpu architecture 'compute_XX'` | Wrong `TORCH_CUDA_ARCH_LIST`; query yours with `python -c "import torch; print(torch.cuda.get_device_capability(0))"` (returns e.g. `(8, 6)` → use `"8.6"`) | +| `g++: internal compiler error: Killed (program cc1plus)` | Out of RAM. Reduce `MAX_JOBS` | +| Build hangs at "Building wheel" with no progress | Patience — first nvcc compile step has no output for ~5 min | + +If the build finishes but +`python -c "from flash_attn import flash_attn_varlen_func"` returns +`None` instead of a function, the install silently failed. Re-check the +log. + +--- + +## sys.path injection — how it actually works + +`MimoAudio` imports look like: + +```python +from src.mimo_audio.mimo_audio import MimoAudio +``` + +So the **parent** of `src/` must be on `sys.path` — i.e., the repo root, +the directory containing `run_mimo_asr.py`. Not `src/mimo_audio/` +itself. The loader does: + +```python +sys.path.insert(0, str(repo_root)) +``` + +Two ways to feed `repo_root` to the loader: + +- **Per-node**: fill in the `repo_path` field on the `MiMoASRLoader` node. +- **Globally**: set `MIMO_ASR_REPO=/path/to/repo` before launching ComfyUI. + Recommended — your workflows then load with `repo_path` blank. + +Symptoms of getting this wrong: + +- `repo_path` blank and env unset → `RuntimeError: MiMo ASR repo path is empty.` +- Path points at `src/mimo_audio/` instead of the repo root → + `RuntimeError: Did not find src/mimo_audio/ under .` +- Path points at a partial clone (no `src/`) → same error. Fix with + `git submodule update --init --recursive`. + +--- + +## Error → Fix table + +Reference card. Each row is a real error message you might see, with the +cause and the fix. + +### Install / configuration errors + +| Symptom | Cause | Fix | +|---|---|---| +| `RuntimeError: MiMo ASR repo path is empty.` | Neither `repo_path` node field nor `MIMO_ASR_REPO` env var is set. | Clone the upstream repo; set the env var or the node field. | +| `RuntimeError: Did not find src/mimo_audio/ under .` | Wrong path, or empty / partial clone. | Verify `ls /src/mimo_audio/`. Run `git submodule update --init --recursive` inside the clone if needed. | +| `RuntimeError: MiMo-V2.5-ASR model path does not exist: ...` | Step 3 skipped, or relative path used. | Download via `hf download`; pass an absolute path to the node. | +| `RuntimeError: MiMo-Audio-Tokenizer path does not exist: ...` | Same as above for the tokenizer. | `hf download XiaomiMiMo/MiMo-Audio-Tokenizer --local-dir ...` | +| `ModuleNotFoundError: No module named 'src'` (at first load) | `sys.path` injection happened but the path is wrong. | Verify `/src/mimo_audio/` exists. | + +### Python dep errors + +| Symptom | Cause | Fix | +|---|---|---| +| `AttributeError: 'MiMoAudioTokenizer' object has no attribute 'all_tied_weights_keys'. Did you mean: '_tied_weights_keys'?` | `transformers` >= 4.50 changed the tied-weights API. MiMo's modeling code uses the old name. | Downgrade: `uv pip install "transformers==4.49.0" "tokenizers>=0.21,<0.22" --force-reinstall` | +| `TypeError: MimoAudio.__init__() got an unexpected keyword argument 'tokenizer_path'` | You're on v0.1.0 of this node pack; it calls `MimoAudio(model_path=..., tokenizer_path=...)`. The actual signature is positional. | Upgrade to v0.1.1+ of this node pack. | +| `ImportError: cannot import name 'XXX' from 'transformers'` | Some other custom node forced `transformers` to a newer or older version after MiMo was working. | Re-run the step 4 install. Consider Strategy B (dedicated ComfyUI). | + +### Runtime errors + +| Symptom | Cause | Fix | +|---|---|---| +| `TypeError: 'NoneType' object is not callable` deep in `mimo_audio_tokenizer/modeling_audio_tokenizer.py` at `flash_attn_varlen_func(...)` | `flash-attn` not installed; upstream binds `flash_attn_varlen_func = None` on `ImportError` and uses it anyway. | Install per Step 5 of the README. There is no SDPA fallback. | +| `RuntimeError: CUDA out of memory` during load | Other models still resident in VRAM. | Restart ComfyUI, or use ComfyUI-Manager's "Free model and node cache". | +| `RuntimeError: FlashAttention only supports fp16 and bf16 dtype` | Loader's dtype detection failed; the model is in fp32. | Open an issue with your `torch.__version__` and the full traceback. | +| Empty transcript output (no error) | `PreviewAny` node missing from your server, but the workflow expects it. | Install rgthree-comfy, or use `ShowText|pysssss`. The CLI has a fallback that reads the asr node's output directly. | +| First run takes 30–60 s | Normal — loading 8B weights from disk. | Subsequent runs in the same ComfyUI session are instant thanks to `_MODEL_CACHE`. | + +--- + +## End-to-end smoke test + +After install, run this from the client side (your dev machine, not the +server): + +```bash +# 0. ComfyUI reachable +curl -s -o /dev/null -w "HTTP %{http_code}\n" http://SERVER:8188/ + +# 1. Nodes registered +curl -s http://SERVER:8188/object_info | python3 -c " +import json, sys +d = json.load(sys.stdin) +for n in ('MiMoASRLoader', 'MiMoASRTranscribe', 'LoadAudio', 'PreviewAny'): + print(f' {n}: {\"OK\" if n in d else \"MISSING\"}')" + +# 2. Dump a workflow without submitting (verifies argparse / paths) +uv run python cli/mimo_asr_cli.py /path/to/test.wav --ip SERVER \ + --model-path /big_disk/models/MiMo-V2.5-ASR \ + --tokenizer-path /big_disk/models/MiMo-Audio-Tokenizer \ + --repo-path /opt/repos/MiMo-V2.5-ASR \ + --dump-workflow /tmp/wf.json --no-run + +# 3. Inspect — load /tmp/wf.json in ComfyUI's UI via "Load" to verify connections + +# 4. Real run +uv run python cli/mimo_asr_cli.py /path/to/test.wav --ip SERVER \ + --model-path /big_disk/models/MiMo-V2.5-ASR \ + --tokenizer-path /big_disk/models/MiMo-Audio-Tokenizer \ + --repo-path /opt/repos/MiMo-V2.5-ASR +``` + +On a 3090, expect: + +- First run: 30–60 s (model load) + 2–4 s (inference) +- Subsequent runs in same session: 2–4 s +- VRAM after load: ~17 GB +- VRAM at peak: ~20–22 GB + +If your numbers are wildly off, file an issue with `nvidia-smi` output +and your `torch` / `transformers` / `flash-attn` versions. diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..a3e4125 --- /dev/null +++ b/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for describing the origin of the Work and + reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Support. While redistributing the Work or + Derivative Works thereof, You may choose to offer, and charge a fee + for, acceptance of support, warranty, indemnity, or other liability + obligations and/or rights consistent with this License. However, + in accepting such obligations, You may act only on Your own behalf + and on Your sole responsibility, not on behalf of any other + Contributor, and only if You agree to indemnify, defend, and hold + each Contributor harmless for any liability incurred by, or claims + asserted against, such Contributor by reason of your accepting any + such warranty or support. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2026 aadebuger + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied. See the License for the specific language governing + permissions and limitations under the License. diff --git a/README.md b/README.md new file mode 100644 index 0000000..2f4cc07 --- /dev/null +++ b/README.md @@ -0,0 +1,290 @@ +# ComfyUI-MiMoASR + +ComfyUI custom nodes for [Xiaomi **MiMo-V2.5-ASR**](https://huggingface.co/XiaomiMiMo/MiMo-V2.5-ASR) +— an 8B end-to-end ASR model with native support for Mandarin-English +code-switching, Chinese dialects (Wu, Cantonese, Hokkien, Sichuanese, +Henan, Northeastern Mandarin), song lyrics, knowledge-heavy speech, and +noisy / multi-speaker audio. Punctuation is generated natively. + +| Node | Purpose | +|---|---| +| 🎯 **MiMo ASR Loader** | Loads `MiMo-V2.5-ASR` (8B) + `MiMo-Audio-Tokenizer`. Cached. | +| 🎯 **MiMo ASR Transcribe** | Runs `asr_sft` on an `AUDIO` input. Outputs `STRING`. | + +--- + +## ⚠️ Read this before installing + +This model is fussier than typical ComfyUI custom nodes. **Three** things +will burn your afternoon if you skip them. The node will install +successfully but fail at run time with confusing errors. + +1. **`transformers` must be `==4.49.0`**. Newer versions changed the + tied-weights API and MiMo's modeling code relies on the old one. + This may conflict with other custom nodes that need `transformers >=4.57` + (Sonic, Qwen3TTS, recent VibeVoice builds). Either accept that or run + MiMo in a dedicated ComfyUI instance. +2. **`flash-attn` is required, not optional.** The audio tokenizer + encoder calls `flash_attn_varlen_func` with **no SDPA fallback**. + Plan to build from source unless your env exactly matches the + upstream wheel (`cu12 + torch 2.6 + Python 3.12`). +3. **You need a local clone of the upstream repo on disk.** + `MimoAudio` is not pip-installable. The loader injects it into + `sys.path` at runtime. + +Quick gut check — verify each row for your env. If any are off, jump +straight to [INSTALL_NOTES.md](./INSTALL_NOTES.md): + +| | Required | Upstream pin | How to check | +|---|---|---|---| +| Python | 3.10+ | 3.12 | `python --version` | +| `torch` | any recent | 2.6.0 | `python -c "import torch; print(torch.__version__)"` | +| `transformers` | **4.49.0** | 4.49.0 | `python -c "import transformers; print(transformers.__version__)"` | +| `flash-attn` | **any working build** | 2.7.4.post1 | `python -c "from flash_attn import flash_attn_varlen_func"` | +| CUDA toolkit | ≥12.0 (for flash-attn build) | 12.x | `nvcc --version` | +| VRAM | 24 GB | — | `nvidia-smi` | + +--- + +## Install (the proven order) + +> **TL;DR** — if you're on Linux with a CUDA toolkit installed and uv-managed +> ComfyUI venv, you can run the one-shot script and skip to Step 6: +> ```bash +> COMFYUI_DIR=/path/to/ComfyUI ./scripts/install_server.sh +> ``` +> It does steps 2–5 below idempotently. Otherwise, follow each step manually. + +Each step has a verification at the end. Don't skip the verifications. + +### 1. Clone this node pack + +```bash +cd ComfyUI/custom_nodes +git clone https://github.com/aadebuger/ComfyUI-MiMoASR.git +``` + +### 2. Clone the upstream repo + +```bash +# Pick any path with disk space; the loader will find it. +git clone https://github.com/XiaomiMiMo/MiMo-V2.5-ASR.git \ + /opt/repos/MiMo-V2.5-ASR + +# Verify: +ls /opt/repos/MiMo-V2.5-ASR/src/mimo_audio/ # must list .py files +``` + +If `src/mimo_audio/` is missing, try `git submodule update --init --recursive` +inside the cloned repo. + +### 3. Download the model weights + +```bash +# Redirect HF cache if your root disk is small. MiMo-V2.5-ASR is ~16-32 GB +# on disk and HF stages a second copy before symlinking. +export HF_HOME=/big_disk/hf_cache +mkdir -p "$HF_HOME" + +pip install -U huggingface-hub + +hf download XiaomiMiMo/MiMo-V2.5-ASR \ + --local-dir /big_disk/models/MiMo-V2.5-ASR + +hf download XiaomiMiMo/MiMo-Audio-Tokenizer \ + --local-dir /big_disk/models/MiMo-Audio-Tokenizer + +# Verify: +ls -lh /big_disk/models/MiMo-V2.5-ASR/*.safetensors +ls -lh /big_disk/models/MiMo-Audio-Tokenizer/ +``` + +### 4. Install Python deps (in ComfyUI's venv) + +```bash +cd ComfyUI && source .venv/bin/activate + +# Snapshot current state — you'll want this if the transformers downgrade +# breaks something else. +uv pip freeze > ~/pip_freeze_before_mimo_$(date +%Y%m%d).txt + +# Pinned transformers + matching tokenizers +uv pip install "transformers==4.49.0" "tokenizers>=0.21,<0.22" --force-reinstall + +# MiMo's runtime deps (from the upstream requirements.txt) +uv pip install "accelerate>=1.9.0" "librosa>=0.11.0" "scipy>=1.16.1" "zhon==2.1.1" + +# Audio I/O for this node pack +uv pip install torchaudio + +# CLI deps (skip if you only use the UI) +uv pip install requests websocket-client + +# Verify: +python -c "import transformers; assert transformers.__version__ == '4.49.0'; print('OK')" +``` + +### 5. Install `flash-attn` + +**Option A — your env matches the upstream wheel** (cu12 + torch 2.6 + py 3.12): + +```bash +uv pip install flash-attn==2.7.4.post1 +``` + +**Option B — build from source** (any newer torch / Python combo). +Budget 20–40 minutes on a 3090 host. + +```bash +uv pip install ninja packaging + +# Adjust to your toolchain +export CUDA_HOME=/usr/local/cuda-12.6 +export TORCH_CUDA_ARCH_LIST="8.6" # 8.6=3090, 8.9=4090, 9.0=H100 +export MAX_JOBS=4 # tune to RAM; lower if you OOM + +# Background-build and tail the log so you can detach +nohup uv pip install flash-attn --no-build-isolation \ + > ~/flash_attn_build.log 2>&1 & +tail -f ~/flash_attn_build.log +``` + +Verify: + +```bash +python -c "from flash_attn import flash_attn_varlen_func; print(flash_attn_varlen_func)" +# Must print: +# NOT: None +``` + +If `flash_attn_varlen_func` is `None`, the install silently failed — +check `~/flash_attn_build.log` for compile errors. See INSTALL_NOTES.md +for the usual suspects (missing `nvcc`, OOM, mismatched +`TORCH_CUDA_ARCH_LIST`). + +### 6. Launch ComfyUI with the env vars + +```bash +export HF_HOME=/big_disk/hf_cache +export MIMO_ASR_REPO=/opt/repos/MiMo-V2.5-ASR + +cd ComfyUI && source .venv/bin/activate +python main.py --listen 0.0.0.0 --port 8188 +``` + +### 7. Verify everything is wired up + +```bash +curl -s http://SERVER:8188/object_info | python3 -c " +import json, sys +d = json.load(sys.stdin) +for n in ('MiMoASRLoader', 'MiMoASRTranscribe', 'LoadAudio', 'PreviewAny'): + print(f' {n}: {\"✓\" if n in d else \"✗\"}')" +``` + +All four should show `✓`. If `PreviewAny` is missing, install +[rgthree-comfy](https://github.com/rgthree/rgthree-comfy) or swap it out +for `ShowText|pysssss` in your workflow. + +--- + +## Usage + +### Minimal workflow + +``` +LoadAudio ──┐ + ├──► 🎯 MiMo ASR Transcribe ──► STRING +MiMo ASR ────┘ +Loader +``` + +**Loader fields** (absolute paths, please): + +- `model_path` — `/big_disk/models/MiMo-V2.5-ASR` +- `tokenizer_path` — `/big_disk/models/MiMo-Audio-Tokenizer` +- `repo_path` — leave blank if `MIMO_ASR_REPO` is set + +**Transcribe fields**: + +- `language` — `Auto` (recommended for code-switched speech), + `Chinese`, or `English` + +The first run loads the 8B model (30–60 s). Subsequent runs in the same +ComfyUI session are instant — the loader has a process-wide cache. + +An API-format example workflow is at +[`example_workflows/mimo_asr_basic.json`](./example_workflows/mimo_asr_basic.json). + +### CLI + +A single-file CLI that uploads an audio file, drives the workflow over +ComfyUI's HTTP + WebSocket API, and prints the transcript: + +```bash +uv run python cli/mimo_asr_cli.py speech.wav \ + --ip SERVER \ + --model-path /big_disk/models/MiMo-V2.5-ASR \ + --tokenizer-path /big_disk/models/MiMo-Audio-Tokenizer \ + --repo-path /opt/repos/MiMo-V2.5-ASR +``` + +See [`cli/README.md`](./cli/README.md) for all flags and a wrapper-alias +recipe to avoid retyping the paths every time. + +--- + +## VRAM + +| | | +|---|---| +| Load (bf16) | ~17 GB | +| Inference peak | ~20–22 GB | +| Minimum GPU | 24 GB (3090, 4090, A100-40G) | +| Speed (30 s audio, 3090) | ~2–4 s after first load | + +A single RTX 3090 is enough. Dual-GPU machines can pin MiMo to one card +via `CUDA_VISIBLE_DEVICES=1` before launching ComfyUI. + +--- + +## Troubleshooting (the five most common) + +For the full table see [INSTALL_NOTES.md](./INSTALL_NOTES.md). + +| Symptom | Cause | Fix | +|---|---|---| +| `RuntimeError: MiMo ASR repo path is empty` | Step 2 missed or env var not set | Clone the repo, set `MIMO_ASR_REPO` | +| `RuntimeError: Did not find src/mimo_audio/` | Empty / partial clone | `git submodule update --init --recursive` inside the clone | +| `RuntimeError: ... model path does not exist` | Step 3 missed, or relative path in node field | Download the weights; use absolute paths | +| `AttributeError: 'MiMoAudioTokenizer' has no attribute 'all_tied_weights_keys'` | `transformers` too new | Downgrade to `4.49.0` | +| `TypeError: 'NoneType' object is not callable` (inside `flash_attn_varlen_func`) | `flash-attn` not installed | Build per Step 5 | + +--- + +## Links + +- Hugging Face: +- Audio tokenizer: +- Upstream repo: +- Blog: +- Live demo: + +--- + +## License + +Apache-2.0 (this node pack). MiMo-V2.5-ASR model weights are MIT-licensed +by Xiaomi; upstream `MimoAudio` source is Apache-2.0. + +## Citation + +```bibtex +@misc{coreteam2026mimov25asr, + title = {MiMo-V2.5-ASR: Robust Speech Recognition Across Languages, + Dialects, and Complex Acoustic Scenarios}, + author = {LLM-Core-Team Xiaomi}, + year = {2026}, + url = {https://github.com/XiaomiMiMo/MiMo-V2.5-ASR} +} +``` diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..3c1ff11 --- /dev/null +++ b/__init__.py @@ -0,0 +1,18 @@ +"""ComfyUI-MiMoASR — Xiaomi MiMo-V2.5-ASR nodes for ComfyUI.""" + +from .nodes.loader import MiMoASRLoader +from .nodes.transcribe import MiMoASRTranscribe + +NODE_CLASS_MAPPINGS = { + "MiMoASRLoader": MiMoASRLoader, + "MiMoASRTranscribe": MiMoASRTranscribe, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "MiMoASRLoader": "🎯 MiMo ASR Loader", + "MiMoASRTranscribe": "🎯 MiMo ASR Transcribe", +} + +WEB_DIRECTORY = None + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/cli/README.md b/cli/README.md new file mode 100644 index 0000000..c543369 --- /dev/null +++ b/cli/README.md @@ -0,0 +1,66 @@ +# ComfyUI-MiMoASR CLI + +Drives the `MiMoASRLoader` + `MiMoASRTranscribe` nodes over ComfyUI's HTTP + +WebSocket API. Same `ComfyClient` template used across the other +`ComfyUI-*` packages. + +## Quick start + +```bash +# uv handles dependencies via the PEP 723 inline script metadata +uv run python mimo_asr_cli.py /path/to/speech.wav --ip 10.8.0.12 + +# Explicit Chinese tag (useful for accented or noisy speech) +uv run python mimo_asr_cli.py speech.wav --language Chinese --ip 10.8.0.12 + +# Pass model + repo paths if your ComfyUI server can't see the defaults +uv run python mimo_asr_cli.py speech.wav \ + --model-path /mnt/models/MiMo-V2.5-ASR \ + --tokenizer-path /mnt/models/MiMo-Audio-Tokenizer \ + --repo-path /mnt/repos/MiMo-V2.5-ASR \ + --ip 10.8.0.12 + +# Save transcript to file +uv run python mimo_asr_cli.py speech.wav --out transcript.txt --ip 10.8.0.12 +``` + +## Flags + +| Flag | Default | Notes | +|---|---|---| +| `audio` (positional) | — | Local audio file. Uploaded to ComfyUI's `input/`. | +| `--ip` | `127.0.0.1` | ComfyUI host. | +| `--port` | `8188` | ComfyUI port. | +| `--server` | — | Full `host:port`, overrides `--ip`/`--port`. | +| `--out` | — | Write transcript to file. | +| `--model-path` | `./models/MiMo-V2.5-ASR` | Relative to ComfyUI cwd, or absolute. | +| `--tokenizer-path` | `./models/MiMo-Audio-Tokenizer` | Same. | +| `--repo-path` | (env `MIMO_ASR_REPO`) | Local clone of [XiaomiMiMo/MiMo-V2.5-ASR](https://github.com/XiaomiMiMo/MiMo-V2.5-ASR). | +| `--language` | `Auto` | `Auto` / `Chinese` / `English`. | +| `--dump-workflow PATH` | — | Save the API-format JSON for inspection. | +| `--no-run` | off | Combine with `--dump-workflow` to skip submission. | +| `--timeout` | `600` | Seconds. | + +## Server-side checklist + +```bash +# 1. ComfyUI is up +curl -s -o /dev/null -w "HTTP %{http_code}\n" http://SERVER:8188/ + +# 2. Nodes loaded +curl -s http://SERVER:8188/object_info | python3 -c " +import json, sys +d = json.load(sys.stdin) +for n in ('MiMoASRLoader', 'MiMoASRTranscribe', 'LoadAudio', 'PreviewAny'): + print(f' {n}: {\"✓\" if n in d else \"✗\"}')" +``` + +`PreviewAny` ships with `rgthree-comfy` (and several others). If it's not +installed, swap the `preview` node in `build_workflow()` for whatever your +environment provides — `ShowText|pysssss` is a drop-in replacement. + +## Audio format notes + +`/upload/image` accepts any file (the endpoint's name is misleading). Audio +gets stored under ComfyUI's `input/` folder, and `LoadAudio` handles +decoding via torchaudio — so wav / mp3 / flac / m4a / ogg all work. diff --git a/cli/mimo_asr_cli.py b/cli/mimo_asr_cli.py new file mode 100644 index 0000000..789b653 --- /dev/null +++ b/cli/mimo_asr_cli.py @@ -0,0 +1,333 @@ +#!/usr/bin/env python3 +# /// script +# requires-python = ">=3.10" +# dependencies = [ +# "requests", +# "websocket-client", +# ] +# /// +"""mimo_asr_cli.py — drive ComfyUI-MiMoASR from the command line. + +Examples +-------- + # 本地音频 → 自动语言检测 + uv run python cli/mimo_asr_cli.py audio.wav --ip 10.8.0.12 + + # 显式指定中文 + uv run python cli/mimo_asr_cli.py speech.flac --language Chinese --ip 10.8.0.12 + + # 自定义模型路径(默认是 ./models/MiMo-V2.5-ASR + ./models/MiMo-Audio-Tokenizer + # 相对于 ComfyUI 启动目录) + uv run python cli/mimo_asr_cli.py speech.wav \ + --model-path /mnt/models/MiMo-V2.5-ASR \ + --tokenizer-path /mnt/models/MiMo-Audio-Tokenizer \ + --repo-path /mnt/repos/MiMo-V2.5-ASR \ + --ip 10.8.0.12 + + # 调试:dump 生成的 workflow 不实际执行 + uv run python cli/mimo_asr_cli.py speech.wav --dump-workflow /tmp/wf.json --no-run +""" + +from __future__ import annotations + +import argparse +import json +import mimetypes +import re +import sys +import time +import urllib.parse +import uuid +from pathlib import Path + +import requests +import websocket # websocket-client + +LANGUAGES = ["Auto", "Chinese", "English"] + + +# --------------------------------------------------------------------------- +# ComfyClient — minimal HTTP + WS wrapper around a ComfyUI server +# --------------------------------------------------------------------------- +class ComfyClient: + def __init__(self, server: str = "127.0.0.1:8188") -> None: + self.server = server + self.http = f"http://{server}" + self.ws = f"ws://{server}/ws" + self.client_id = str(uuid.uuid4()) + + def upload_image(self, file_path: Path, *, overwrite: bool = True) -> str: + """Upload any file via POST /upload/image. Works for audio too.""" + mime, _ = mimetypes.guess_type(str(file_path)) + with open(file_path, "rb") as fh: + files = {"image": (file_path.name, fh, mime or "application/octet-stream")} + data = {"overwrite": "true" if overwrite else "false"} + r = requests.post(f"{self.http}/upload/image", files=files, data=data, timeout=60) + r.raise_for_status() + payload = r.json() + # ComfyUI returns {"name": "", ...} + return payload["name"] + + def object_info(self) -> dict: + r = requests.get(f"{self.http}/object_info", timeout=15) + r.raise_for_status() + return r.json() + + def submit(self, workflow: dict) -> str: + r = requests.post( + f"{self.http}/prompt", + json={"prompt": workflow, "client_id": self.client_id}, + timeout=30, + ) + r.raise_for_status() + payload = r.json() + if "prompt_id" not in payload: + raise RuntimeError(f"Submit failed: {payload}") + return payload["prompt_id"] + + def history(self, prompt_id: str) -> dict: + r = requests.get(f"{self.http}/history/{prompt_id}", timeout=15) + r.raise_for_status() + data = r.json() + return data.get(prompt_id, {}) + + def view(self, filename: str, subfolder: str, type_: str) -> bytes: + params = urllib.parse.urlencode( + {"filename": filename, "subfolder": subfolder, "type": type_} + ) + r = requests.get(f"{self.http}/view?{params}", timeout=60) + r.raise_for_status() + return r.content + + def run(self, workflow: dict, *, timeout_s: float = 600.0) -> dict: + """Submit a workflow and stream progress until completion.""" + prompt_id = self.submit(workflow) + print(f" prompt_id = {prompt_id}") + + ws = websocket.create_connection(f"{self.ws}?clientId={self.client_id}", timeout=30) + ws.settimeout(timeout_s) + deadline = time.time() + timeout_s + last_node = None + try: + while True: + if time.time() > deadline: + raise TimeoutError(f"Workflow exceeded {timeout_s}s") + msg = ws.recv() + if isinstance(msg, (bytes, bytearray)): + continue + try: + payload = json.loads(msg) + except json.JSONDecodeError: + continue + t = payload.get("type") + d = payload.get("data", {}) + if d.get("prompt_id") and d.get("prompt_id") != prompt_id: + continue + if t == "progress": + cur, tot = d.get("value", 0), d.get("max", 1) + sys.stdout.write(f"\r ▸ {cur:>4}/{tot}") + sys.stdout.flush() + elif t == "executing": + node = d.get("node") + if node is None: + # execution finished for this prompt + sys.stdout.write("\n") + break + if node != last_node: + sys.stdout.write(f"\r ▸ node {str(node):<20}\n") + sys.stdout.flush() + last_node = node + elif t == "execution_error": + sys.stdout.write("\n") + raise RuntimeError( + f"Workflow error in node {d.get('node_id')}: " + f"{d.get('exception_message')}" + ) + finally: + try: + ws.close() + except Exception: # noqa: BLE001 + pass + return self.history(prompt_id) + + +# --------------------------------------------------------------------------- +# Workflow construction +# --------------------------------------------------------------------------- +def build_workflow( + *, + audio_filename: str, + model_path: str, + tokenizer_path: str, + repo_path: str, + language: str, +) -> dict: + """Return the API-format workflow dict.""" + return { + "loader": { + "class_type": "MiMoASRLoader", + "inputs": { + "model_path": model_path, + "tokenizer_path": tokenizer_path, + "repo_path": repo_path, + "force_reload": False, + }, + }, + "load_audio": { + "class_type": "LoadAudio", + "inputs": { + "audio": audio_filename, + }, + }, + "asr": { + "class_type": "MiMoASRTranscribe", + "inputs": { + "model": ["loader", 0], + "audio": ["load_audio", 0], + "language": language, + }, + }, + # Capture STRING output so we can read it back from /history. + "preview": { + "class_type": "PreviewAny", + "inputs": { + "source": ["asr", 0], + }, + }, + } + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- +def extract_text_from_history(history: dict, node_id: str = "preview") -> str: + """Pull the text payload PreviewAny stored under outputs[node_id].""" + outs = history.get("outputs", {}) + if node_id not in outs: + # Fall back to the asr node itself if PreviewAny was missing. + for nid, node_out in outs.items(): + if "text" in node_out: + node_id = nid + break + else: + return "" + node_out = outs[node_id] + if "text" in node_out: + v = node_out["text"] + return "\n".join(str(x) for x in v) if isinstance(v, list) else str(v) + parts: list[str] = [] + for v in node_out.values(): + if isinstance(v, list): + parts.extend(str(x) for x in v) + return "\n".join(parts) + + +_SAFE_NAME_RE = re.compile(r"[^A-Za-z0-9._-]+") + + +def sanitize_filename(name: str) -> str: + sanitized = _SAFE_NAME_RE.sub("_", name).strip("._-") + return sanitized or "audio" + + +# --------------------------------------------------------------------------- +# main +# --------------------------------------------------------------------------- +def main() -> int: + ap = argparse.ArgumentParser(description="ComfyUI-MiMoASR CLI") + ap.add_argument("audio", type=Path, help="Local audio file to transcribe (wav/mp3/flac/m4a/...)") + + # Server. + ap.add_argument("--ip", default="127.0.0.1") + ap.add_argument("--port", type=int, default=8188) + ap.add_argument("--server", default=None, help="Full host:port, overrides --ip/--port") + + # Output. + ap.add_argument("--out", type=Path, default=None, + help="Write transcript to file (default: stdout only).") + + # MiMo paths. + ap.add_argument("--model-path", default="./models/MiMo-V2.5-ASR", + help="Path to MiMo-V2.5-ASR weights (relative to ComfyUI cwd or absolute).") + ap.add_argument("--tokenizer-path", default="./models/MiMo-Audio-Tokenizer", + help="Path to MiMo-Audio-Tokenizer.") + ap.add_argument("--repo-path", default="", + help="Path to local clone of XiaomiMiMo/MiMo-V2.5-ASR. " + "Leave blank to use the MIMO_ASR_REPO env var on the server.") + + # Generation. + ap.add_argument("--language", choices=LANGUAGES, default="Auto") + + # Debug. + ap.add_argument("--dump-workflow", type=Path, default=None, + help="Save the generated workflow JSON for inspection.") + ap.add_argument("--no-run", action="store_true", + help="Skip submission (use with --dump-workflow to just produce the JSON).") + ap.add_argument("--timeout", type=float, default=600.0) + + args = ap.parse_args() + + server = args.server or f"{args.ip}:{args.port}" + if not args.audio.is_file(): + print(f"Error: audio file not found: {args.audio}", file=sys.stderr) + return 2 + + client = ComfyClient(server=server) + + # 1) Upload audio. + if not args.no_run: + print(f"▸ Uploading {args.audio.name} to {server} ...") + uploaded_name = client.upload_image(args.audio) + print(f" uploaded as: {uploaded_name}") + else: + uploaded_name = sanitize_filename(args.audio.name) + + # 2) Build workflow. + workflow = build_workflow( + audio_filename=uploaded_name, + model_path=args.model_path, + tokenizer_path=args.tokenizer_path, + repo_path=args.repo_path, + language=args.language, + ) + + if args.dump_workflow: + args.dump_workflow.parent.mkdir(parents=True, exist_ok=True) + args.dump_workflow.write_text( + json.dumps(workflow, indent=2, ensure_ascii=False), encoding="utf-8" + ) + print(f"▸ Workflow JSON written to {args.dump_workflow}") + + if args.no_run: + print("▸ --no-run set, exiting without submission.") + return 0 + + # 3) Run. + print(f"▸ Submitting workflow (language={args.language!r}) ...") + t0 = time.time() + history = client.run(workflow, timeout_s=args.timeout) + elapsed = time.time() - t0 + print(f"▸ Workflow finished in {elapsed:.2f}s") + + # 4) Extract text. + text = extract_text_from_history(history, node_id="preview") + if not text: + print("Warning: no text output found in workflow history.", file=sys.stderr) + print(json.dumps(history.get("outputs", {}), indent=2, ensure_ascii=False)[:2000], + file=sys.stderr) + return 1 + + print("\n----- Transcript -----") + print(text) + print("----------------------") + + if args.out: + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(text + "\n", encoding="utf-8") + print(f"▸ Saved to {args.out}") + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/example_workflows/mimo_asr_basic.json b/example_workflows/mimo_asr_basic.json new file mode 100644 index 0000000..ab18117 --- /dev/null +++ b/example_workflows/mimo_asr_basic.json @@ -0,0 +1,31 @@ +{ + "loader": { + "class_type": "MiMoASRLoader", + "inputs": { + "model_path": "./models/MiMo-V2.5-ASR", + "tokenizer_path": "./models/MiMo-Audio-Tokenizer", + "repo_path": "", + "force_reload": false + } + }, + "load_audio": { + "class_type": "LoadAudio", + "inputs": { + "audio": "sample.wav" + } + }, + "asr": { + "class_type": "MiMoASRTranscribe", + "inputs": { + "model": ["loader", 0], + "audio": ["load_audio", 0], + "language": "Auto" + } + }, + "preview": { + "class_type": "PreviewAny", + "inputs": { + "source": ["asr", 0] + } + } +} diff --git a/nodes/__init__.py b/nodes/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/nodes/_types.py b/nodes/_types.py new file mode 100644 index 0000000..21f655c --- /dev/null +++ b/nodes/_types.py @@ -0,0 +1,18 @@ +"""Shared types and constants for ComfyUI-MiMoASR.""" + +# ComfyUI custom type name for the MiMo model bundle returned by the loader. +MIMO_ASR_MODEL = "MIMO_ASR_MODEL" + +# Language tags accepted by MimoAudio.asr_sft(audio_tag=...). +LANGUAGE_TAGS = { + "Auto": "", + "Chinese": "", + "English": "", +} + +# Node category shown in the ComfyUI search bar. +CATEGORY = "🎯 MiMo ASR" + +# Env var the loader falls back to when `repo_path` is left blank in the node UI. +# Should point at a local clone of https://github.com/XiaomiMiMo/MiMo-V2.5-ASR +ENV_REPO_PATH = "MIMO_ASR_REPO" diff --git a/nodes/loader.py b/nodes/loader.py new file mode 100644 index 0000000..c11c05a --- /dev/null +++ b/nodes/loader.py @@ -0,0 +1,155 @@ +"""MiMoASRLoader — load the MiMo-V2.5-ASR model + audio tokenizer.""" + +from __future__ import annotations + +import os +import sys +from pathlib import Path + +from ._types import CATEGORY, ENV_REPO_PATH, MIMO_ASR_MODEL + +# Process-wide cache so reruns of the same workflow don't reload the 8B model. +_MODEL_CACHE: dict[str, dict] = {} + + +def _resolve_repo_path(repo_path: str) -> str: + """Return the resolved repo path, or raise with a useful error.""" + candidate = repo_path.strip() or os.environ.get(ENV_REPO_PATH, "").strip() + if not candidate: + raise RuntimeError( + "MiMo ASR repo path is empty. Either:\n" + " 1. Set the `repo_path` field on the MiMoASRLoader node, OR\n" + f" 2. Set the {ENV_REPO_PATH} env var before launching ComfyUI.\n\n" + "The path must point to a local clone of:\n" + " https://github.com/XiaomiMiMo/MiMo-V2.5-ASR\n" + "i.e. the directory that contains the `src/mimo_audio/` folder." + ) + + repo = Path(candidate).expanduser().resolve() + src_dir = repo / "src" / "mimo_audio" + if not src_dir.is_dir(): + raise RuntimeError( + f"Did not find `src/mimo_audio/` under {repo}.\n" + "Make sure repo_path points at the root of the MiMo-V2.5-ASR repo " + "(the folder containing `run_mimo_asr.py` and the `src/` directory)." + ) + + repo_str = str(repo) + if repo_str not in sys.path: + sys.path.insert(0, repo_str) + return repo_str + + +class MiMoASRLoader: + """Load Xiaomi MiMo-V2.5-ASR (8B) and its audio tokenizer. + + The MimoAudio class is not pip-installable, so this node sources it + from a local clone of the upstream repo. Loading takes ~30-60 s on a + 3090 and consumes ~16 GB of VRAM — the result is cached by + (model_path, tokenizer_path, repo_path) so subsequent nodes reuse it. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model_path": ( + "STRING", + { + "default": "./models/MiMo-V2.5-ASR", + "multiline": False, + "tooltip": "Local path to the MiMo-V2.5-ASR weights " + "(downloaded from XiaomiMiMo/MiMo-V2.5-ASR).", + }, + ), + "tokenizer_path": ( + "STRING", + { + "default": "./models/MiMo-Audio-Tokenizer", + "multiline": False, + "tooltip": "Local path to the MiMo-Audio-Tokenizer " + "(downloaded from XiaomiMiMo/MiMo-Audio-Tokenizer).", + }, + ), + "repo_path": ( + "STRING", + { + "default": "", + "multiline": False, + "tooltip": f"Path to a clone of XiaomiMiMo/MiMo-V2.5-ASR " + f"(the folder containing src/mimo_audio/). " + f"Leave blank to read from the {ENV_REPO_PATH} env var.", + }, + ), + }, + "optional": { + "force_reload": ( + "BOOLEAN", + {"default": False, "tooltip": "Bypass the in-memory model cache."}, + ), + }, + } + + RETURN_TYPES = (MIMO_ASR_MODEL,) + RETURN_NAMES = ("model",) + FUNCTION = "load" + CATEGORY = CATEGORY + + def load( + self, + model_path: str, + tokenizer_path: str, + repo_path: str, + force_reload: bool = False, + ): + model_path = str(Path(model_path).expanduser()) + tokenizer_path = str(Path(tokenizer_path).expanduser()) + repo = _resolve_repo_path(repo_path) + + cache_key = f"{model_path}|{tokenizer_path}|{repo}" + if not force_reload and cache_key in _MODEL_CACHE: + return (_MODEL_CACHE[cache_key],) + + # Lazy import — only after sys.path is patched and only when loading + # is actually requested. Avoids importing torch/transformers at + # ComfyUI startup. + from src.mimo_audio.mimo_audio import MimoAudio # type: ignore + + if not Path(model_path).is_dir(): + raise RuntimeError( + f"MiMo-V2.5-ASR model path does not exist: {model_path}\n" + "Download with:\n" + " hf download XiaomiMiMo/MiMo-V2.5-ASR " + f"--local-dir {model_path}" + ) + if not Path(tokenizer_path).is_dir(): + raise RuntimeError( + f"MiMo-Audio-Tokenizer path does not exist: {tokenizer_path}\n" + "Download with:\n" + " hf download XiaomiMiMo/MiMo-Audio-Tokenizer " + f"--local-dir {tokenizer_path}" + ) + + print(f"[MiMoASRLoader] Loading MiMo-V2.5-ASR from {model_path}") + print(f"[MiMoASRLoader] Loading audio tokenizer from {tokenizer_path}") + # Upstream's run_mimo_asr.py uses positional args. The HF README shows + # kwargs that don't match the actual __init__ signature — call by + # position to stay compatible with both old and new upstream commits. + model = MimoAudio(model_path, tokenizer_path) + + bundle = { + "model": model, + "model_path": model_path, + "tokenizer_path": tokenizer_path, + "repo_path": repo, + } + _MODEL_CACHE[cache_key] = bundle + print("[MiMoASRLoader] Model ready.") + return (bundle,) + + @classmethod + def IS_CHANGED(cls, model_path, tokenizer_path, repo_path, force_reload=False): + # When force_reload is on, return a sentinel that always differs. + if force_reload: + return float("NaN") + return f"{model_path}|{tokenizer_path}|{repo_path}" diff --git a/nodes/transcribe.py b/nodes/transcribe.py new file mode 100644 index 0000000..c785b67 --- /dev/null +++ b/nodes/transcribe.py @@ -0,0 +1,123 @@ +"""MiMoASRTranscribe — run MiMo-V2.5-ASR on an audio input.""" + +from __future__ import annotations + +import os +import tempfile +from pathlib import Path + +from ._types import CATEGORY, LANGUAGE_TAGS, MIMO_ASR_MODEL + + +def _audio_dict_to_wav(audio: dict, out_path: Path) -> None: + """Persist a ComfyUI AUDIO dict to a wav file on disk. + + AUDIO is `{"waveform": Tensor[B, C, N], "sample_rate": int}`. + We squeeze the batch axis and pick channel 0 if multi-channel — the + upstream MimoAudio pipeline expects a mono 16k-style waveform but + handles its own resampling internally. + """ + import torch # noqa: F401 (kept local — heavy import) + import torchaudio + + waveform = audio["waveform"] + sample_rate = int(audio["sample_rate"]) + + if waveform.ndim == 3: + # [B, C, N] → [C, N] (take first item in batch) + waveform = waveform[0] + elif waveform.ndim == 1: + waveform = waveform.unsqueeze(0) + + # MimoAudio expects 16 kHz mono internally; downmix to mono here and + # leave resampling to the model's own preprocessing. + if waveform.size(0) > 1: + waveform = waveform.mean(dim=0, keepdim=True) + + # torchaudio.save needs a float CPU tensor. + waveform = waveform.detach().to("cpu").float() + torchaudio.save(str(out_path), waveform, sample_rate) + + +class MiMoASRTranscribe: + """Transcribe an AUDIO input using a loaded MiMo-V2.5-ASR model. + + Supports automatic language detection (recommended for code-switched + speech) or explicit Chinese / English language tags. Punctuation is + produced natively by the model — no post-processing needed. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": (MIMO_ASR_MODEL,), + "audio": ("AUDIO",), + "language": ( + list(LANGUAGE_TAGS.keys()), + { + "default": "Auto", + "tooltip": "Auto = automatic language detection " + "(recommended for code-switched speech).", + }, + ), + }, + "optional": { + "keep_temp_file": ( + "BOOLEAN", + { + "default": False, + "tooltip": "If True, keep the intermediate wav file " + "(useful for debugging).", + }, + ), + }, + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("text",) + FUNCTION = "transcribe" + CATEGORY = CATEGORY + OUTPUT_NODE = False + + def transcribe( + self, + model: dict, + audio: dict, + language: str = "Auto", + keep_temp_file: bool = False, + ): + if not isinstance(model, dict) or "model" not in model: + raise RuntimeError( + "MiMoASRTranscribe got an invalid model bundle. " + "Connect the `model` output of MiMoASRLoader." + ) + if not isinstance(audio, dict) or "waveform" not in audio: + raise RuntimeError( + "MiMoASRTranscribe expected a ComfyUI AUDIO dict on `audio` " + "(connect a LoadAudio node)." + ) + + audio_tag = LANGUAGE_TAGS.get(language, "") + mimo = model["model"] + + # MimoAudio.asr_sft takes a file path, so dump the waveform to a tempfile. + fd, tmp_path = tempfile.mkstemp(prefix="mimoasr_", suffix=".wav") + os.close(fd) + tmp = Path(tmp_path) + + try: + _audio_dict_to_wav(audio, tmp) + print(f"[MiMoASRTranscribe] Transcribing {tmp} (lang={language!r})") + text = mimo.asr_sft(str(tmp), audio_tag=audio_tag) + if text is None: + text = "" + text = str(text).strip() + print(f"[MiMoASRTranscribe] {len(text)} chars produced.") + return (text,) + finally: + if not keep_temp_file: + try: + tmp.unlink(missing_ok=True) + except Exception: # noqa: BLE001 + pass diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..d9c2fc5 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,44 @@ +[build-system] +requires = ["setuptools>=68", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "comfyui-mimoasr" +version = "0.1.2" +description = "ComfyUI custom nodes for Xiaomi MiMo-V2.5-ASR speech recognition" +readme = "README.md" +license = { text = "Apache-2.0" } +requires-python = ">=3.10" +authors = [ + { name = "aadebuger" }, +] +keywords = ["comfyui", "asr", "speech-recognition", "mimo", "xiaomi", "whisper"] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "License :: OSI Approved :: Apache Software License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Topic :: Multimedia :: Sound/Audio :: Speech", + "Topic :: Scientific/Engineering :: Artificial Intelligence", +] +dependencies = [ + "torchaudio", +] + +[project.optional-dependencies] +cli = [ + "requests", + "websocket-client", +] + +[project.urls] +Homepage = "https://github.com/aadebuger/ComfyUI-MiMoASR" +Issues = "https://github.com/aadebuger/ComfyUI-MiMoASR/issues" +Upstream = "https://github.com/XiaomiMiMo/MiMo-V2.5-ASR" +HuggingFace = "https://huggingface.co/XiaomiMiMo/MiMo-V2.5-ASR" + +[tool.setuptools] +packages = ["nodes"] diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..d19db36 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,9 @@ +# Runtime deps for the ComfyUI node pack itself. +# (Heavy deps — torch, transformers, accelerate, librosa, flash-attn — +# must come from the upstream MiMo-V2.5-ASR repo's requirements.txt, +# installed into your ComfyUI venv. See INSTALL_NOTES.md.) +torchaudio + +# CLI deps (only needed for cli/mimo_asr_cli.py) +requests +websocket-client diff --git a/scripts/install_server.sh b/scripts/install_server.sh new file mode 100755 index 0000000..3e09f9c --- /dev/null +++ b/scripts/install_server.sh @@ -0,0 +1,194 @@ +#!/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 <