315 lines
10 KiB
Python
315 lines
10 KiB
Python
#!/usr/bin/env python3
|
|
"""Lens-Turbo CLI — submit a T2I workflow to a running ComfyUI server.
|
|
|
|
Pattern A (prompt + args) from the v2 README.
|
|
|
|
Usage:
|
|
python cli/lens_turbo_cli.py "a cat holding a hello-world sign" \\
|
|
--width 1024 --height 1024 \\
|
|
--base-resolution 1440 --aspect-ratio 1:1 \\
|
|
--steps 4 --cfg 1.0 --seed 0 \\
|
|
--ip 10.8.0.12 --gpu 0 \\
|
|
--out ./out
|
|
|
|
The CLI is deliberately small: it builds the API-format workflow, submits
|
|
it, listens for completion over WebSocket (with the §10 socket-level recv
|
|
timeout!), and saves the output PNG.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import io
|
|
import json
|
|
import sys
|
|
import time
|
|
import urllib.parse
|
|
import urllib.request
|
|
import uuid
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
try:
|
|
import websocket # type: ignore # `uv pip install websocket-client`
|
|
except ImportError:
|
|
print(
|
|
"Missing dependency: pip install websocket-client",
|
|
file=sys.stderr,
|
|
)
|
|
raise
|
|
|
|
|
|
class ComfyClient:
|
|
def __init__(
|
|
self,
|
|
server: str = "127.0.0.1:8188",
|
|
*,
|
|
ws_recv_timeout: float = 600.0,
|
|
connect_timeout: float = 30.0,
|
|
) -> None:
|
|
self.server = server
|
|
self.client_id = str(uuid.uuid4())
|
|
self.ws_recv_timeout = ws_recv_timeout
|
|
self.connect_timeout = connect_timeout
|
|
|
|
def _http(self, path: str, data: bytes | None = None, method: str | None = None):
|
|
url = f"http://{self.server}{path}"
|
|
req = urllib.request.Request(url, data=data, method=method)
|
|
if data is not None:
|
|
req.add_header("Content-Type", "application/json")
|
|
with urllib.request.urlopen(req, timeout=self.connect_timeout) as r:
|
|
return r.read()
|
|
|
|
def submit(self, workflow: dict) -> str:
|
|
payload = json.dumps({"prompt": workflow, "client_id": self.client_id})
|
|
raw = self._http("/prompt", data=payload.encode("utf-8"), method="POST")
|
|
return json.loads(raw)["prompt_id"]
|
|
|
|
def history(self, prompt_id: str) -> dict:
|
|
raw = self._http(f"/history/{prompt_id}")
|
|
all_history = json.loads(raw)
|
|
return all_history.get(prompt_id, {})
|
|
|
|
def view(self, filename: str, subfolder: str, type_: str) -> bytes:
|
|
qs = urllib.parse.urlencode(
|
|
{"filename": filename, "subfolder": subfolder, "type": type_}
|
|
)
|
|
return self._http(f"/view?{qs}")
|
|
|
|
def run(self, workflow: dict, *, timeout_s: float = 1800.0) -> dict:
|
|
prompt_id = self.submit(workflow)
|
|
print(f"[cli] prompt_id={prompt_id}", flush=True)
|
|
|
|
ws = websocket.WebSocket()
|
|
ws.connect(
|
|
f"ws://{self.server}/ws?clientId={self.client_id}",
|
|
timeout=self.connect_timeout,
|
|
)
|
|
# ⭐ Critical: socket-level recv timeout. Cold-loading GPT-OSS-20B
|
|
# can take 5+ minutes with zero WS traffic.
|
|
ws.settimeout(self.ws_recv_timeout)
|
|
|
|
deadline = time.time() + timeout_s
|
|
done = False
|
|
try:
|
|
while time.time() < deadline:
|
|
try:
|
|
raw = ws.recv()
|
|
except websocket.WebSocketTimeoutException:
|
|
# Fall back to /history polling.
|
|
h = self.history(prompt_id)
|
|
status = h.get("status", {})
|
|
if status.get("completed") or status.get("status_str") in (
|
|
"success",
|
|
"error",
|
|
):
|
|
done = True
|
|
break
|
|
print(
|
|
"[cli] WS idle, /history says not done — continuing",
|
|
flush=True,
|
|
)
|
|
continue
|
|
if isinstance(raw, bytes):
|
|
continue # binary frames are previews
|
|
msg = json.loads(raw)
|
|
mtype = msg.get("type")
|
|
data = msg.get("data", {})
|
|
if mtype == "executing" and data.get("prompt_id") == prompt_id:
|
|
if data.get("node") is None:
|
|
done = True
|
|
break
|
|
elif mtype == "execution_error":
|
|
raise RuntimeError(f"server error: {data}")
|
|
finally:
|
|
try:
|
|
ws.close()
|
|
except Exception:
|
|
pass
|
|
|
|
if not done:
|
|
raise TimeoutError(f"timed out after {timeout_s}s waiting on prompt")
|
|
|
|
return self.history(prompt_id)
|
|
|
|
|
|
def build_workflow(
|
|
*,
|
|
prompt: str,
|
|
negative_prompt: str,
|
|
base_resolution: int,
|
|
aspect_ratio: str,
|
|
steps: int,
|
|
cfg: float,
|
|
seed: int,
|
|
device: str,
|
|
dtype: str,
|
|
offload_mode: str,
|
|
cpu_construct: bool,
|
|
repo_id: str,
|
|
force_reload: bool,
|
|
) -> dict:
|
|
"""Build API-format workflow (string node ids, inputs as [src, idx])."""
|
|
return {
|
|
"loader": {
|
|
"class_type": "LensTurboLoader",
|
|
"inputs": {
|
|
"repo_id": repo_id,
|
|
"device": device,
|
|
"dtype": dtype,
|
|
"offload_mode": offload_mode,
|
|
"cpu_construct": cpu_construct,
|
|
"force_reload": force_reload,
|
|
"evict_others": True,
|
|
"hf_cache_dir": "",
|
|
},
|
|
},
|
|
"sampler": {
|
|
"class_type": "LensTurboSampler",
|
|
"inputs": {
|
|
"pipeline": ["loader", 0],
|
|
"prompt": prompt,
|
|
"negative_prompt": negative_prompt,
|
|
"base_resolution": base_resolution,
|
|
"aspect_ratio": aspect_ratio,
|
|
"num_inference_steps": steps,
|
|
"guidance_scale": cfg,
|
|
"seed": seed,
|
|
},
|
|
},
|
|
"save": {
|
|
"class_type": "SaveImage",
|
|
"inputs": {
|
|
"images": ["sampler", 0],
|
|
"filename_prefix": "LensTurbo",
|
|
},
|
|
},
|
|
}
|
|
|
|
|
|
def main() -> int:
|
|
ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
|
|
|
|
ap.add_argument("prompt", help="positive prompt text")
|
|
ap.add_argument("--negative-prompt", default="")
|
|
|
|
# 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="host:port (overrides ip/port)")
|
|
|
|
# GPU
|
|
ap.add_argument("--gpu", type=int, default=0, help="cuda:N for the loader")
|
|
|
|
# Model selection
|
|
ap.add_argument("--repo-id", default="microsoft/Lens-Turbo")
|
|
ap.add_argument(
|
|
"--dtype",
|
|
default="bf16",
|
|
choices=["bf16", "fp16", "fp32"],
|
|
)
|
|
ap.add_argument(
|
|
"--offload-mode",
|
|
default="cpu_text_encoder",
|
|
choices=["cpu_text_encoder", "sequential", "model", "none"],
|
|
help=(
|
|
"VRAM strategy. cpu_text_encoder=hybrid (default, ~11 GB peak, "
|
|
"24 GB safe), sequential=broken with Lens, model=sub-model "
|
|
"(≥48 GB), none=no offload (≥48 GB)."
|
|
),
|
|
)
|
|
ap.add_argument(
|
|
"--no-cpu-construct",
|
|
action="store_true",
|
|
help="Skip the `with torch.device('cpu')` construct guard",
|
|
)
|
|
ap.add_argument("--force-reload", action="store_true")
|
|
|
|
# Sampling
|
|
ap.add_argument("--base-resolution", type=int, default=1440)
|
|
ap.add_argument(
|
|
"--aspect-ratio",
|
|
default="1:1",
|
|
choices=[
|
|
"1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3", "21:9", "9:21",
|
|
],
|
|
)
|
|
ap.add_argument("--steps", type=int, default=4)
|
|
ap.add_argument("--cfg", type=float, default=1.0)
|
|
ap.add_argument("--seed", type=int, default=0)
|
|
|
|
# I/O
|
|
ap.add_argument("--out", type=Path, default=Path("./out"))
|
|
ap.add_argument(
|
|
"--dump-workflow",
|
|
type=Path,
|
|
default=None,
|
|
help="Save the generated workflow JSON (debug, load in UI).",
|
|
)
|
|
|
|
# Timeouts
|
|
ap.add_argument("--timeout", type=float, default=1800.0)
|
|
ap.add_argument(
|
|
"--ws-recv-timeout",
|
|
type=float,
|
|
default=600.0,
|
|
help="Socket-level idle timeout. First cold load can take 5+ minutes.",
|
|
)
|
|
|
|
args = ap.parse_args()
|
|
server = args.server or f"{args.ip}:{args.port}"
|
|
device = "cpu" if args.gpu < 0 else f"cuda:{args.gpu}"
|
|
|
|
wf = build_workflow(
|
|
prompt=args.prompt,
|
|
negative_prompt=args.negative_prompt,
|
|
base_resolution=args.base_resolution,
|
|
aspect_ratio=args.aspect_ratio,
|
|
steps=args.steps,
|
|
cfg=args.cfg,
|
|
seed=args.seed,
|
|
device=device,
|
|
dtype=args.dtype,
|
|
offload_mode=args.offload_mode,
|
|
cpu_construct=not args.no_cpu_construct,
|
|
repo_id=args.repo_id,
|
|
force_reload=args.force_reload,
|
|
)
|
|
|
|
if args.dump_workflow is not None:
|
|
args.dump_workflow.parent.mkdir(parents=True, exist_ok=True)
|
|
args.dump_workflow.write_text(json.dumps(wf, indent=2, ensure_ascii=False))
|
|
print(f"[cli] workflow dumped to {args.dump_workflow}", flush=True)
|
|
|
|
args.out.mkdir(parents=True, exist_ok=True)
|
|
|
|
client = ComfyClient(server, ws_recv_timeout=args.ws_recv_timeout)
|
|
print(f"[cli] submitting to {server} ...", flush=True)
|
|
t0 = time.time()
|
|
history = client.run(wf, timeout_s=args.timeout)
|
|
dt = time.time() - t0
|
|
print(f"[cli] complete in {dt:.1f}s", flush=True)
|
|
|
|
# Pull images
|
|
outputs = history.get("outputs", {})
|
|
save_out = outputs.get("save", {})
|
|
images = save_out.get("images", [])
|
|
if not images:
|
|
print(
|
|
f"[cli] no images in history.outputs['save']! full outputs:\n"
|
|
f"{json.dumps(outputs, indent=2, default=str)[:2000]}",
|
|
file=sys.stderr,
|
|
)
|
|
return 2
|
|
|
|
for i, meta in enumerate(images):
|
|
raw = client.view(meta["filename"], meta.get("subfolder", ""), meta.get("type", "output"))
|
|
target = args.out / f"lens_turbo_{args.seed}_{i:02d}.png"
|
|
target.write_bytes(raw)
|
|
print(f"[cli] saved {target}", flush=True)
|
|
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|