Patrick Plate
134 lines
4.4 KiB
Python
134 lines
4.4 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Quick CLI for generating images via ComfyUI + FLUX.2 Klein Heretic.
|
||
|
||
Usage:
|
||
python gen.py "your prompt here"
|
||
python gen.py "your prompt" --steps 20 --width 1280 --height 720
|
||
python gen.py "your prompt" --seed 12345
|
||
python gen.py "your prompt" --count 3
|
||
|
||
Output saved to ~/Pictures/mcp-generated/
|
||
"""
|
||
|
||
import argparse
|
||
import json
|
||
import random
|
||
import sys
|
||
import time
|
||
import urllib.request
|
||
from pathlib import Path
|
||
|
||
COMFYUI = "http://localhost:8188"
|
||
OUTPUT_DIR = Path.home() / "Pictures" / "mcp-generated"
|
||
WORKFLOW_PATH = Path(__file__).parent / "src/workflows/flux2_klein_heretic.json"
|
||
|
||
|
||
def load_workflow():
|
||
with open(WORKFLOW_PATH) as f:
|
||
return json.load(f)
|
||
|
||
|
||
def submit(workflow):
|
||
data = json.dumps({"prompt": workflow}).encode()
|
||
req = urllib.request.Request(
|
||
f"{COMFYUI}/prompt", data=data,
|
||
headers={"Content-Type": "application/json"}
|
||
)
|
||
with urllib.request.urlopen(req) as resp:
|
||
return json.loads(resp.read())["prompt_id"]
|
||
|
||
|
||
def wait(prompt_id, timeout=300):
|
||
print(" ⏳ Waiting for ComfyUI...", end="", flush=True)
|
||
start = time.time()
|
||
while time.time() - start < timeout:
|
||
with urllib.request.urlopen(f"{COMFYUI}/history/{prompt_id}") as resp:
|
||
history = json.loads(resp.read())
|
||
if prompt_id in history:
|
||
print(" done.", flush=True)
|
||
outputs = history[prompt_id].get("outputs", {})
|
||
for node_out in outputs.values():
|
||
if "images" in node_out:
|
||
return node_out["images"][0]
|
||
return None
|
||
print(".", end="", flush=True)
|
||
time.sleep(2)
|
||
raise TimeoutError(f"Timed out after {timeout}s")
|
||
|
||
|
||
def download(filename, subfolder=""):
|
||
url = f"{COMFYUI}/view?filename={filename}&subfolder={subfolder}&type=output"
|
||
with urllib.request.urlopen(url) as resp:
|
||
return resp.read()
|
||
|
||
|
||
def generate(prompt, steps=20, width=1024, height=1024, seed=-1, name="cli"):
|
||
if seed == -1:
|
||
seed = random.randint(0, 2**32 - 1)
|
||
|
||
workflow = load_workflow()
|
||
|
||
# Patch positive prompt (node 2)
|
||
workflow["2"]["inputs"]["text"] = prompt
|
||
# Patch negative prompt (node 3) — leave empty
|
||
workflow["3"]["inputs"]["text"] = ""
|
||
# Patch seed (node 10)
|
||
if "10" in workflow:
|
||
workflow["10"]["inputs"]["noise_seed"] = seed
|
||
# Patch dimensions (node 6)
|
||
workflow["6"]["inputs"]["width"] = width
|
||
workflow["6"]["inputs"]["height"] = height
|
||
# Patch steps (node 7)
|
||
workflow["7"]["inputs"]["steps"] = steps
|
||
# Patch output filename (node 13)
|
||
workflow["13"]["inputs"]["filename_prefix"] = name
|
||
|
||
print(f" Prompt : {prompt[:80]}{'...' if len(prompt) > 80 else ''}")
|
||
print(f" Size : {width}×{height} Steps: {steps} Seed: {seed}")
|
||
|
||
prompt_id = submit(workflow)
|
||
image_info = wait(prompt_id)
|
||
|
||
if not image_info:
|
||
print(" ❌ No output image returned.", file=sys.stderr)
|
||
return None
|
||
|
||
img_data = download(image_info["filename"], image_info.get("subfolder", ""))
|
||
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
|
||
out_path = OUTPUT_DIR / f"{name}_{seed}.png"
|
||
out_path.write_bytes(img_data)
|
||
print(f" ✅ Saved: {out_path} ({len(img_data) // 1024}KB)")
|
||
return out_path
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(
|
||
description="Generate images via ComfyUI FLUX.2 Klein Heretic"
|
||
)
|
||
parser.add_argument("prompt", help="Text prompt for the image")
|
||
parser.add_argument("--steps", type=int, default=20, help="Inference steps (default: 20)")
|
||
parser.add_argument("--width", type=int, default=1024, help="Width in pixels (default: 1024)")
|
||
parser.add_argument("--height", type=int, default=1024, help="Height in pixels (default: 1024)")
|
||
parser.add_argument("--seed", type=int, default=-1, help="Seed (-1 = random)")
|
||
parser.add_argument("--count", type=int, default=1, help="Number of images (default: 1)")
|
||
parser.add_argument("--name", default="cli", help="Output filename prefix (default: cli)")
|
||
args = parser.parse_args()
|
||
|
||
for i in range(args.count):
|
||
if args.count > 1:
|
||
print(f"\n[{i+1}/{args.count}]")
|
||
seed = args.seed if args.seed != -1 else -1
|
||
generate(
|
||
prompt=args.prompt,
|
||
steps=args.steps,
|
||
width=args.width,
|
||
height=args.height,
|
||
seed=seed,
|
||
name=f"{args.name}_{i+1:02d}" if args.count > 1 else args.name,
|
||
)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|