334 lines
12 KiB
Python
334 lines
12 KiB
Python
#!/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": "<server-side filename>", ...}
|
|
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())
|