423 lines
16 KiB
Python
Executable File
423 lines
16 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""hymt2_cli.py — single-shot Hy-MT2 translation via ComfyUI.
|
|
|
|
Usage:
|
|
python hymt2_cli.py "今天天气真好。" --target English --ip 10.8.0.12
|
|
python hymt2_cli.py - --target 日语 --prompt-lang zh --ip 10.8.0.12 # stdin
|
|
python hymt2_cli.py path/to/text.txt --target French --ip 10.8.0.12
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
import urllib.parse
|
|
import urllib.request
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
import websocket # pip install websocket-client
|
|
|
|
|
|
# Use /mnt/modelall/tmp if available (v2.0 §1: /tmp is only 10 GB on the server).
|
|
_PREFERRED_TMP = "/mnt/modelall/tmp"
|
|
if os.path.isdir(_PREFERRED_TMP):
|
|
tempfile.tempdir = _PREFERRED_TMP
|
|
|
|
|
|
# ----------------------------------------------------------------------------
|
|
# WebSocket completion signals
|
|
# ----------------------------------------------------------------------------
|
|
# Modern ComfyUI (the version Hua deployed in 2026-05) emits a different set of
|
|
# WS messages than the older docs describe. From actual /history capture:
|
|
# prompt_id: ...
|
|
# status messages:
|
|
# - execution_start
|
|
# - execution_cached (some nodes hit the cache)
|
|
# - execution_success (whole prompt finished)
|
|
#
|
|
# Crucially, in fully-cached or very fast paths, the older `executing` /
|
|
# `node=null` "queue empty" message is NOT emitted at all, so a client that
|
|
# only listens for that signal will block on ws.recv() forever.
|
|
#
|
|
# v0.1.1 handles all of:
|
|
# - execution_success -> our prompt completed cleanly
|
|
# - execution_error -> our prompt failed
|
|
# - executing + node=None -> legacy completion signal (still emitted on the
|
|
# non-cached path)
|
|
# - executed -> single-node output (informational, keep listening)
|
|
# - execution_cached -> informational only
|
|
# - progress -> progress bar
|
|
|
|
_COMPLETION_TYPES = {"execution_success"}
|
|
_ERROR_TYPES = {"execution_error"}
|
|
|
|
|
|
class ComfyClient:
|
|
"""Minimal ComfyUI client.
|
|
|
|
v0.1.2: opens WS BEFORE submit + polls /history on status messages to
|
|
avoid the cached-prompt race that would block on ws.recv().
|
|
v0.1.1: handles modern execution_success WS messages.
|
|
"""
|
|
|
|
def __init__(self, server="127.0.0.1:8188", *,
|
|
ws_recv_timeout=600.0, connect_timeout=30.0, verbose=False):
|
|
self.server = server
|
|
self.client_id = str(uuid.uuid4())
|
|
self.ws_recv_timeout = ws_recv_timeout
|
|
self.connect_timeout = connect_timeout
|
|
self.verbose = verbose
|
|
|
|
def submit(self, workflow):
|
|
payload = json.dumps({"prompt": workflow, "client_id": self.client_id}).encode()
|
|
req = urllib.request.Request(
|
|
f"http://{self.server}/prompt",
|
|
data=payload,
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
with urllib.request.urlopen(req, timeout=self.connect_timeout) as r:
|
|
data = json.loads(r.read())
|
|
return data["prompt_id"]
|
|
|
|
def history(self, prompt_id):
|
|
url = f"http://{self.server}/history/{urllib.parse.quote(prompt_id)}"
|
|
with urllib.request.urlopen(url, timeout=self.connect_timeout) as r:
|
|
data = json.loads(r.read())
|
|
return data.get(prompt_id, {})
|
|
|
|
def _log(self, msg):
|
|
if self.verbose:
|
|
print(f"[hymt2/cli] {msg}", file=sys.stderr, flush=True)
|
|
|
|
def run(self, workflow, *, timeout_s=1800):
|
|
# v0.1.2: open WS BEFORE submit to avoid the cached-prompt race.
|
|
# If we submit first then connect, the server may broadcast all
|
|
# execution_start / execution_success events before our WS is up,
|
|
# and we'd block on ws.recv() forever waiting for events that already
|
|
# happened. Opening WS first means the server's broadcast lands in our
|
|
# socket buffer.
|
|
ws = websocket.WebSocket()
|
|
ws.connect(
|
|
f"ws://{self.server}/ws?clientId={self.client_id}",
|
|
timeout=self.connect_timeout,
|
|
)
|
|
# ⭐ v2.0 §10: ws.connect(timeout) is connect-only. Set the recv timeout
|
|
# so cold load (model download / first-load) doesn't trip the socket.
|
|
ws.settimeout(self.ws_recv_timeout)
|
|
|
|
try:
|
|
prompt_id = self.submit(workflow)
|
|
self._log(f"submitted prompt_id={prompt_id}")
|
|
|
|
# v0.1.2: Even with WS-first, an instant cache hit can complete
|
|
# during the submit() POST round trip. Poll /history once before
|
|
# entering the WS loop.
|
|
h0 = self.history(prompt_id)
|
|
if h0 and h0.get("status", {}).get("completed"):
|
|
self._log("prompt already completed at submit time (cache hit)")
|
|
return h0
|
|
|
|
deadline = time.time() + timeout_s
|
|
completed = False
|
|
last_exec_err = None
|
|
|
|
while not completed:
|
|
if time.time() > deadline:
|
|
raise TimeoutError(f"workflow exceeded {timeout_s}s overall")
|
|
|
|
try:
|
|
raw = ws.recv()
|
|
except (websocket.WebSocketTimeoutException, TimeoutError):
|
|
# Idle timeout — poll /history as fallback. Server may have
|
|
# completed and we just missed the WS signal.
|
|
self._log("ws recv idle timeout, polling /history")
|
|
h = self.history(prompt_id)
|
|
if h and h.get("status", {}).get("completed"):
|
|
self._log("fallback /history poll found completion")
|
|
return h
|
|
continue
|
|
except websocket.WebSocketConnectionClosedException:
|
|
# Server closed the socket. Fall through to /history.
|
|
self._log("ws connection closed, polling /history")
|
|
break
|
|
|
|
if isinstance(raw, bytes):
|
|
# Binary frames are progress images we don't care about.
|
|
continue
|
|
try:
|
|
msg = json.loads(raw)
|
|
except Exception:
|
|
continue
|
|
|
|
kind = msg.get("type")
|
|
data = msg.get("data", {}) or {}
|
|
msg_pid = data.get("prompt_id")
|
|
|
|
if self.verbose:
|
|
s = json.dumps(data, ensure_ascii=False, default=str)
|
|
if len(s) > 200:
|
|
s = s[:200] + "...<truncated>"
|
|
self._log(f"ws msg type={kind} data={s}")
|
|
|
|
# --- Completion: execution_success for our prompt ---
|
|
if kind in _COMPLETION_TYPES and msg_pid == prompt_id:
|
|
completed = True
|
|
break
|
|
|
|
# --- Error: execution_error for our prompt ---
|
|
if kind in _ERROR_TYPES and msg_pid == prompt_id:
|
|
last_exec_err = data
|
|
break
|
|
|
|
# --- Legacy completion signal: executing/node=null ---
|
|
if kind == "executing" and data.get("node") is None:
|
|
# Some older versions don't set prompt_id; accept either
|
|
# match or missing.
|
|
if msg_pid is None or msg_pid == prompt_id:
|
|
completed = True
|
|
break
|
|
|
|
# --- Progress bar ---
|
|
if kind == "progress":
|
|
val = data.get("value")
|
|
mx = data.get("max")
|
|
if val is not None and mx:
|
|
print(f" progress {val}/{mx}",
|
|
file=sys.stderr, flush=True)
|
|
|
|
# v0.1.2: status with queue_remaining=0 is a strong hint that
|
|
# the server already finished. Cross-check /history.
|
|
if kind == "status":
|
|
qrem = (data.get("status", {})
|
|
.get("exec_info", {})
|
|
.get("queue_remaining"))
|
|
if qrem == 0:
|
|
h = self.history(prompt_id)
|
|
if h and h.get("status", {}).get("completed"):
|
|
self._log("status/queue_remaining=0 + history "
|
|
"shows completed -> race detected")
|
|
return h
|
|
|
|
# Other types (execution_start, execution_cached, executed)
|
|
# are informational; we keep listening.
|
|
finally:
|
|
try:
|
|
ws.close()
|
|
except Exception:
|
|
pass
|
|
|
|
if last_exec_err:
|
|
raise RuntimeError(
|
|
f"ComfyUI execution_error: "
|
|
f"{json.dumps(last_exec_err, ensure_ascii=False)[:1000]}"
|
|
)
|
|
|
|
# ALWAYS fetch /history at the end. Outputs are only guaranteed to be
|
|
# in /history; the WebSocket success message is just a signal.
|
|
h = self.history(prompt_id)
|
|
if self.verbose:
|
|
status = h.get("status", {})
|
|
self._log(f"final /history: completed={status.get('completed')} "
|
|
f"status_str={status.get('status_str')}")
|
|
return h
|
|
|
|
|
|
def extract_text_from_history(history, node_id) -> str:
|
|
"""v2.0 §5.3 — peel STRING outputs out of /history."""
|
|
outs = history.get("outputs", {})
|
|
if node_id not in outs:
|
|
return ""
|
|
no = outs[node_id]
|
|
if "text" in no:
|
|
v = no["text"]
|
|
return "\n".join(str(x) for x in v) if isinstance(v, list) else str(v)
|
|
parts = []
|
|
for v in no.values():
|
|
if isinstance(v, list):
|
|
parts.extend(str(x) for x in v)
|
|
return "\n".join(parts)
|
|
|
|
|
|
def build_workflow(*, model_path, device, dtype, source_text, target_lang,
|
|
prompt_lang, temperature, top_p, top_k, repetition_penalty,
|
|
max_new_tokens, seed, terminology, target_style,
|
|
background_text, force_reload):
|
|
"""Return the API-format ComfyUI workflow dict (v2.0 §5.2 style)."""
|
|
return {
|
|
"loader": {
|
|
"class_type": "HyMT2Loader",
|
|
"inputs": {
|
|
"model_path": model_path,
|
|
"device": device,
|
|
"dtype": dtype,
|
|
"force_reload": force_reload,
|
|
"evict_others": True,
|
|
},
|
|
},
|
|
"translate": {
|
|
"class_type": "HyMT2Translate",
|
|
"inputs": {
|
|
"model": ["loader", 0],
|
|
"source_text": source_text,
|
|
"target_lang": target_lang,
|
|
"prompt_lang": prompt_lang,
|
|
"temperature": temperature,
|
|
"top_p": top_p,
|
|
"top_k": top_k,
|
|
"repetition_penalty": repetition_penalty,
|
|
"max_new_tokens": max_new_tokens,
|
|
"seed": seed,
|
|
"terminology": terminology,
|
|
"target_style": target_style,
|
|
"background_text": background_text,
|
|
},
|
|
},
|
|
"preview": {
|
|
"class_type": "PreviewAny",
|
|
"inputs": {"source": ["translate", 0]},
|
|
},
|
|
}
|
|
|
|
|
|
def read_source(src_arg: str) -> str:
|
|
"""Resolve CLI arg into actual text:
|
|
- '-' → stdin
|
|
- a path that exists → file contents
|
|
- else → treat as inline text
|
|
"""
|
|
if src_arg == "-":
|
|
return sys.stdin.read()
|
|
p = Path(src_arg)
|
|
if p.exists() and p.is_file():
|
|
return p.read_text(encoding="utf-8")
|
|
return src_arg
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(description="Hy-MT2 translation via ComfyUI.")
|
|
ap.add_argument("source",
|
|
help="Source text. Use '-' for stdin, or a file path, "
|
|
"or inline string.")
|
|
|
|
# Server (v2.0 §6)
|
|
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.")
|
|
|
|
# Model selection
|
|
ap.add_argument("--model-path", default="tencent/Hy-MT2-1.8B",
|
|
help="HF repo id or absolute local dir.")
|
|
ap.add_argument("--gpu", type=int, default=0,
|
|
help="CUDA device index → cuda:N")
|
|
ap.add_argument("--device", default=None,
|
|
help="Override --gpu with an explicit device string.")
|
|
ap.add_argument("--dtype", default="bf16", choices=["bf16", "fp16", "fp32"])
|
|
|
|
# Translation params
|
|
ap.add_argument("--target", default="English",
|
|
help="Target language NAME. Use English names with "
|
|
"--prompt-lang en, Chinese names with --prompt-lang zh.")
|
|
ap.add_argument("--prompt-lang", default="en", choices=["en", "zh"])
|
|
ap.add_argument("--temperature", type=float, default=0.7)
|
|
ap.add_argument("--top-p", type=float, default=0.6)
|
|
ap.add_argument("--top-k", type=int, default=20)
|
|
ap.add_argument("--repetition-penalty", type=float, default=1.05)
|
|
ap.add_argument("--max-new-tokens", type=int, default=4096)
|
|
ap.add_argument("--seed", type=int, default=-1)
|
|
|
|
# Optional template extras
|
|
ap.add_argument("--terminology", default="",
|
|
help="Path to a glossary file (multi-line, 'src=>tgt' format) "
|
|
"or the inline glossary string.")
|
|
ap.add_argument("--style", default="", help="Target style descriptor.")
|
|
ap.add_argument("--background", default="",
|
|
help="Background paragraph or path to a file.")
|
|
|
|
# Output / debug (v2.0 §6 standard set)
|
|
ap.add_argument("--out", type=Path, default=None,
|
|
help="Write the translation to this file. "
|
|
"If not set, print to stdout.")
|
|
ap.add_argument("--dump-workflow", type=Path, default=None,
|
|
help="Dump the generated API workflow JSON for UI loading.")
|
|
ap.add_argument("--timeout", type=float, default=1800.0)
|
|
ap.add_argument("--ws-recv-timeout", type=float, default=600.0)
|
|
ap.add_argument("--force-reload", action="store_true",
|
|
help="Skip the loader cache (force re-load).")
|
|
ap.add_argument("-v", "--verbose", action="store_true",
|
|
help="Print every received WebSocket message and the "
|
|
"final /history status.")
|
|
|
|
args = ap.parse_args()
|
|
|
|
server = args.server or f"{args.ip}:{args.port}"
|
|
device = args.device or f"cuda:{args.gpu}"
|
|
|
|
source_text = read_source(args.source)
|
|
|
|
def _resolve_optional_text(value: str) -> str:
|
|
if not value:
|
|
return ""
|
|
p = Path(value)
|
|
if p.exists() and p.is_file():
|
|
return p.read_text(encoding="utf-8")
|
|
return value
|
|
|
|
terminology_text = _resolve_optional_text(args.terminology)
|
|
background_text = _resolve_optional_text(args.background)
|
|
|
|
wf = build_workflow(
|
|
model_path=args.model_path,
|
|
device=device,
|
|
dtype=args.dtype,
|
|
source_text=source_text,
|
|
target_lang=args.target,
|
|
prompt_lang=args.prompt_lang,
|
|
temperature=args.temperature,
|
|
top_p=args.top_p,
|
|
top_k=args.top_k,
|
|
repetition_penalty=args.repetition_penalty,
|
|
max_new_tokens=args.max_new_tokens,
|
|
seed=args.seed,
|
|
terminology=terminology_text,
|
|
target_style=args.style,
|
|
background_text=background_text,
|
|
force_reload=args.force_reload,
|
|
)
|
|
|
|
if args.dump_workflow:
|
|
args.dump_workflow.parent.mkdir(parents=True, exist_ok=True)
|
|
args.dump_workflow.write_text(json.dumps(wf, ensure_ascii=False, indent=2))
|
|
print(f"workflow dumped to {args.dump_workflow}", file=sys.stderr)
|
|
|
|
client = ComfyClient(
|
|
server=server,
|
|
ws_recv_timeout=args.ws_recv_timeout,
|
|
verbose=args.verbose,
|
|
)
|
|
|
|
t0 = time.time()
|
|
history = client.run(wf, timeout_s=args.timeout)
|
|
dt = time.time() - t0
|
|
|
|
text = extract_text_from_history(history, "preview")
|
|
if not text:
|
|
text = extract_text_from_history(history, "translate")
|
|
|
|
print(f"[hymt2/cli] done in {dt:.1f}s", file=sys.stderr)
|
|
|
|
if args.out:
|
|
args.out.parent.mkdir(parents=True, exist_ok=True)
|
|
args.out.write_text(text, encoding="utf-8")
|
|
print(f"[hymt2/cli] wrote {len(text)} chars to {args.out}", file=sys.stderr)
|
|
else:
|
|
print(text)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|