317 lines
10 KiB
Python
317 lines
10 KiB
Python
#!/usr/bin/env -S uv run --script
|
|
# /// script
|
|
# requires-python = ">=3.10"
|
|
# dependencies = [
|
|
# "websocket-client>=1.7",
|
|
# ]
|
|
# ///
|
|
"""English vocabulary card CLI.
|
|
|
|
End-to-end: prompt -> SDXL/Z-Image -> Florence-2 detection -> annotator -> save.
|
|
|
|
Examples
|
|
--------
|
|
uv run english_card_cli.py \\
|
|
"flat vector illustration of a tidy classroom: desk, chair, blackboard, books, pencil, backpack, clock, globe, computer" \\
|
|
--ip 10.8.0.12 --steps 25
|
|
|
|
uv run english_card_cli.py \\
|
|
"a sunny kitchen scene, flat vector illustration" \\
|
|
--words "apple,banana,knife,cutting board,bowl,kettle,faucet,refrigerator" \\
|
|
--ip 10.8.0.12
|
|
"""
|
|
import argparse
|
|
import json
|
|
import sys
|
|
import time
|
|
import urllib.parse
|
|
import urllib.request
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
import websocket # websocket-client
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# ComfyClient — minimal HTTP + WebSocket client for the ComfyUI server.
|
|
# ---------------------------------------------------------------------------
|
|
class ComfyClient:
|
|
def __init__(self, server: str = "127.0.0.1:8188"):
|
|
self.server = server
|
|
self.client_id = str(uuid.uuid4())
|
|
|
|
def submit(self, workflow: dict) -> str:
|
|
body = json.dumps(
|
|
{"prompt": workflow, "client_id": self.client_id}
|
|
).encode()
|
|
req = urllib.request.Request(
|
|
f"http://{self.server}/prompt",
|
|
data=body,
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
return json.loads(urllib.request.urlopen(req).read())["prompt_id"]
|
|
|
|
def history(self, pid: str) -> dict:
|
|
r = urllib.request.urlopen(f"http://{self.server}/history/{pid}")
|
|
return json.loads(r.read()).get(pid, {})
|
|
|
|
def view(self, filename: str, subfolder: str, type_: str) -> bytes:
|
|
q = urllib.parse.urlencode(
|
|
{"filename": filename, "subfolder": subfolder, "type": type_}
|
|
)
|
|
return urllib.request.urlopen(f"http://{self.server}/view?{q}").read()
|
|
|
|
def run(self, workflow: dict, timeout_s: float = 600.0) -> dict:
|
|
ws = websocket.WebSocket()
|
|
ws.connect(f"ws://{self.server}/ws?clientId={self.client_id}")
|
|
pid = self.submit(workflow)
|
|
start, last = time.time(), None
|
|
try:
|
|
while time.time() - start < timeout_s:
|
|
msg = ws.recv()
|
|
if not isinstance(msg, str):
|
|
continue
|
|
m = json.loads(msg)
|
|
t, d = m.get("type"), m.get("data", {})
|
|
|
|
if t == "progress":
|
|
cur, tot = d.get("value", 0), d.get("max", 1)
|
|
print(f"\r ▸ {cur:>4}/{tot}", end="", flush=True)
|
|
|
|
elif t == "executing":
|
|
if d.get("prompt_id") == pid and d.get("node") is None:
|
|
print() # newline after progress
|
|
return self.history(pid)
|
|
node = d.get("node")
|
|
if node and node != last:
|
|
last = node
|
|
print(f"\r ▸ node {node:<22}", end="", flush=True)
|
|
|
|
elif t == "execution_error":
|
|
raise RuntimeError(
|
|
f"Node {d.get('node_id')}: {d.get('exception_message')}"
|
|
)
|
|
finally:
|
|
ws.close()
|
|
raise TimeoutError(f"timeout after {timeout_s}s")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Workflow builder — returns API-format dict.
|
|
# ---------------------------------------------------------------------------
|
|
def build_workflow(
|
|
prompt: str,
|
|
*,
|
|
ckpt: str,
|
|
width: int,
|
|
height: int,
|
|
steps: int,
|
|
cfg: float,
|
|
seed: int,
|
|
target_words: list | None = None,
|
|
font_size: int = 36,
|
|
box_width: int = 4,
|
|
florence2_model: str = "microsoft/Florence-2-large",
|
|
florence2_precision: str = "fp16",
|
|
) -> dict:
|
|
"""SDXL/Z-Image -> Florence-2 -> Annotator -> Save."""
|
|
if target_words:
|
|
f2_task = "caption_to_phrase_grounding"
|
|
f2_text = ", ".join(target_words)
|
|
else:
|
|
f2_task = "dense_region_caption"
|
|
f2_text = ""
|
|
|
|
return {
|
|
"loader": {
|
|
"class_type": "CheckpointLoaderSimple",
|
|
"inputs": {"ckpt_name": ckpt},
|
|
},
|
|
"pos": {
|
|
"class_type": "CLIPTextEncode",
|
|
"inputs": {"text": prompt, "clip": ["loader", 1]},
|
|
},
|
|
"neg": {
|
|
"class_type": "CLIPTextEncode",
|
|
"inputs": {
|
|
"text": "blurry, low quality, watermark, text overlay",
|
|
"clip": ["loader", 1],
|
|
},
|
|
},
|
|
"latent": {
|
|
"class_type": "EmptySD3LatentImage",
|
|
"inputs": {"width": width, "height": height, "batch_size": 1},
|
|
},
|
|
"sampler": {
|
|
"class_type": "KSampler",
|
|
"inputs": {
|
|
"model": ["loader", 0],
|
|
"seed": seed,
|
|
"steps": steps,
|
|
"cfg": cfg,
|
|
"sampler_name": "euler",
|
|
"scheduler": "simple",
|
|
"denoise": 1.0,
|
|
"positive": ["pos", 0],
|
|
"negative": ["neg", 0],
|
|
"latent_image": ["latent", 0],
|
|
},
|
|
},
|
|
"decode": {
|
|
"class_type": "VAEDecode",
|
|
"inputs": {"samples": ["sampler", 0], "vae": ["loader", 2]},
|
|
},
|
|
"f2_load": {
|
|
"class_type": "Florence2ModelLoader",
|
|
"inputs": {
|
|
"model": florence2_model,
|
|
"precision": florence2_precision,
|
|
"attention": "sdpa",
|
|
},
|
|
},
|
|
"f2_run": {
|
|
"class_type": "Florence2Run",
|
|
"inputs": {
|
|
"image": ["decode", 0],
|
|
"florence2_model": ["f2_load", 0],
|
|
"task": f2_task,
|
|
"text_input": f2_text,
|
|
"fill_mask": False,
|
|
"keep_model_loaded": True,
|
|
"max_new_tokens": 1024,
|
|
"num_beams": 3,
|
|
"do_sample": False,
|
|
"output_mask_select": "",
|
|
},
|
|
},
|
|
"annotate": {
|
|
"class_type": "EnglishWordAnnotator",
|
|
"inputs": {
|
|
"image": ["decode", 0],
|
|
# Florence2Run outputs vary across versions; the annotator
|
|
# accepts any shape. Adjust indices if your version differs.
|
|
"bboxes": ["f2_run", 2],
|
|
"labels": ["f2_run", 3],
|
|
"font_size": font_size,
|
|
"box_width": box_width,
|
|
"label_position": "above",
|
|
},
|
|
},
|
|
"save": {
|
|
"class_type": "SaveImage",
|
|
"inputs": {
|
|
"images": ["annotate", 0],
|
|
"filename_prefix": "english_card",
|
|
},
|
|
},
|
|
}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Entry point.
|
|
# ---------------------------------------------------------------------------
|
|
def main():
|
|
ap = argparse.ArgumentParser(
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
description=__doc__,
|
|
)
|
|
ap.add_argument("prompt", help="scene prompt")
|
|
ap.add_argument(
|
|
"--words",
|
|
help="comma-separated target words; uses Florence-2 phrase grounding",
|
|
)
|
|
|
|
# generation params
|
|
ap.add_argument("--ckpt", default="sd_xl_base_1.0.safetensors")
|
|
ap.add_argument("--width", type=int, default=1024)
|
|
ap.add_argument("--height", type=int, default=1024)
|
|
ap.add_argument("--steps", type=int, default=25)
|
|
ap.add_argument("--cfg", type=float, default=6.0)
|
|
ap.add_argument(
|
|
"--seed", type=int, default=int(time.time()) & 0xFFFFFFFF
|
|
)
|
|
|
|
# florence-2
|
|
ap.add_argument(
|
|
"--florence2-model", default="microsoft/Florence-2-large"
|
|
)
|
|
ap.add_argument(
|
|
"--florence2-precision", default="fp16",
|
|
choices=["fp16", "bf16", "fp32"],
|
|
)
|
|
|
|
# annotation
|
|
ap.add_argument("--font-size", type=int, default=36)
|
|
ap.add_argument("--box-width", type=int, default=4)
|
|
|
|
# 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")
|
|
|
|
# output / debug
|
|
ap.add_argument("--out", type=Path, default=Path("./out"))
|
|
ap.add_argument(
|
|
"--dump-workflow",
|
|
type=Path,
|
|
default=None,
|
|
help="save generated workflow JSON (debug)",
|
|
)
|
|
ap.add_argument("--timeout", type=float, default=600)
|
|
|
|
args = ap.parse_args()
|
|
server = args.server or f"{args.ip}:{args.port}"
|
|
args.out.mkdir(parents=True, exist_ok=True)
|
|
|
|
target_words = (
|
|
[w.strip() for w in args.words.split(",") if w.strip()]
|
|
if args.words
|
|
else None
|
|
)
|
|
|
|
wf = build_workflow(
|
|
args.prompt,
|
|
ckpt=args.ckpt,
|
|
width=args.width,
|
|
height=args.height,
|
|
steps=args.steps,
|
|
cfg=args.cfg,
|
|
seed=args.seed,
|
|
target_words=target_words,
|
|
font_size=args.font_size,
|
|
box_width=args.box_width,
|
|
florence2_model=args.florence2_model,
|
|
florence2_precision=args.florence2_precision,
|
|
)
|
|
|
|
if args.dump_workflow:
|
|
args.dump_workflow.write_text(json.dumps(wf, indent=2))
|
|
print(f"workflow dumped -> {args.dump_workflow}")
|
|
|
|
print(f"server: {server} seed: {args.seed}")
|
|
if target_words:
|
|
print(f"target words: {', '.join(target_words)}")
|
|
|
|
client = ComfyClient(server)
|
|
hist = client.run(wf, timeout_s=args.timeout)
|
|
|
|
imgs = hist.get("outputs", {}).get("save", {}).get("images", [])
|
|
if not imgs:
|
|
print("no images in history. outputs:")
|
|
print(json.dumps(hist.get("outputs", {}), indent=2, default=str)[:1200])
|
|
sys.exit(1)
|
|
|
|
for img in imgs:
|
|
data = client.view(
|
|
img["filename"],
|
|
img.get("subfolder", ""),
|
|
img.get("type", "output"),
|
|
)
|
|
dst = args.out / img["filename"]
|
|
dst.write_bytes(data)
|
|
print(f"saved -> {dst} ({len(data) // 1024} KB)")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|