nodes/aadebuger-ComfyUI-EnglishAnnotator
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Initial commit: ComfyUI-EnglishAnnotator

Browse Source
This commit is contained in:
hua zhuang
2026-05-20 13:01:58 +08:00
commit c81663f9a2
14 changed files with 1078 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cli/README.md
+82
View File
@@ -0,0 +1,82 @@
# CLI: english_card_cli.py
End-to-end CLI for the English vocabulary card pipeline:
**text prompt → generated scene → Florence-2 detection → annotated image**.
Uses ComfyUI's WebSocket API. PEP 723 inline dependencies, so `uv run` sets up the venv automatically.
## Run
Free detection (Florence-2 picks what it sees):
```bash
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
```
Fixed vocabulary (uses `caption_to_phrase_grounding` — most reliable for daily card pipelines):
```bash
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
```
## Arguments
| Arg | Default | Purpose |
|---|---|---|
| `prompt` | (required) | Scene prompt |
| `--words` | — | Comma-separated target words; switches Florence-2 to `caption_to_phrase_grounding` |
| `--ckpt` | `sd_xl_base_1.0.safetensors` | Checkpoint filename in `models/checkpoints/` |
| `--width` / `--height` | 1024 / 1024 | Output dimensions |
| `--steps` | 25 | Sampler steps |
| `--cfg` | 6.0 | CFG scale |
| `--seed` | (timestamp) | Sampler seed |
| `--florence2-model` | `microsoft/Florence-2-large` | HF model id |
| `--florence2-precision` | `fp16` | `fp16` / `bf16` / `fp32` |
| `--font-size` | 36 | Label text size |
| `--box-width` | 4 | Box outline width |
| `--ip` / `--port` | `127.0.0.1` / `8188` | ComfyUI server |
| `--server` | — | Full `host:port` (overrides `--ip`/`--port`) |
| `--out` | `./out` | Output directory |
| `--dump-workflow` | — | Save generated API workflow JSON for debugging |
| `--timeout` | 600 | Seconds before giving up |
## Debugging
Dump the workflow JSON, then inspect or replay:
```bash
./english_card_cli.py "scene prompt" --dump-workflow /tmp/wf.json
# 1. Load /tmp/wf.json in the ComfyUI UI (Workflow → Load (API Format)) to inspect
# 2. Or POST it directly:
curl -X POST -H "Content-Type: application/json" \
-d "$(jq -n --slurpfile p /tmp/wf.json '{prompt: $p[0]}')" \
http://10.8.0.12:8188/prompt
```
Inspect the latest history (to see actual node output shapes):
```bash
curl -s http://10.8.0.12:8188/history | python3 -c "
import json, sys
h = json.load(sys.stdin)
last = sorted(h)[-1]
print(json.dumps(h[last]['outputs'], indent=2, default=str)[:1500])
"
```
## Common adjustments
- **BBOX index mismatch** — if Florence-2 returns bboxes at a different output index, edit the
`annotate` node in `build_workflow()`:
```python
"bboxes": ["f2_run", 2], # try 1, 2, or 3 depending on Florence2 version
"labels": ["f2_run", 3],
```
- **Z-Image-Turbo** — set `--steps 4 --cfg 1.0` and use the turbo checkpoint name.
- **Flux** — adjust `sampler_name` / `scheduler` in `build_workflow()` accordingly.
cli/english_card_cli.py
+316
View File
@@ -0,0 +1,316 @@
#!/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()