228 lines
9.0 KiB
Python
228 lines
9.0 KiB
Python
"""
|
|
Shared utilities for all eval scripts.
|
|
ComfyUI API, file discovery, image loading, progress logging.
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import glob
|
|
import json
|
|
import time
|
|
import uuid
|
|
import urllib.request
|
|
import urllib.error
|
|
|
|
COMFY_URL = "http://localhost:8188"
|
|
|
|
|
|
# ── ComfyUI API ──────────────────────────────────────────────
|
|
|
|
def check_comfyui():
|
|
try:
|
|
resp = urllib.request.urlopen(f"{COMFY_URL}/system_stats", timeout=5)
|
|
stats = json.loads(resp.read())
|
|
gpu = stats.get("devices", [{}])[0].get("name", "unknown")
|
|
print(f"[OK] ComfyUI connected ({gpu})")
|
|
return True
|
|
except Exception:
|
|
print(f"[ERROR] Cannot connect to ComfyUI at {COMFY_URL}")
|
|
print(" Start ComfyUI first.")
|
|
sys.exit(1)
|
|
|
|
|
|
def queue_prompt(workflow):
|
|
payload = json.dumps({
|
|
"prompt": workflow,
|
|
"client_id": str(uuid.uuid4()),
|
|
}).encode("utf-8")
|
|
req = urllib.request.Request(
|
|
f"{COMFY_URL}/prompt",
|
|
data=payload,
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
resp = urllib.request.urlopen(req)
|
|
return json.loads(resp.read())["prompt_id"]
|
|
|
|
|
|
def wait_for_completion(prompt_id, timeout=600):
|
|
start = time.time()
|
|
while time.time() - start < timeout:
|
|
try:
|
|
resp = urllib.request.urlopen(f"{COMFY_URL}/history/{prompt_id}")
|
|
history = json.loads(resp.read())
|
|
if prompt_id in history:
|
|
status = history[prompt_id].get("status", {})
|
|
if status.get("status_str") == "error":
|
|
return False
|
|
if history[prompt_id].get("outputs"):
|
|
return True
|
|
except urllib.error.URLError:
|
|
pass
|
|
time.sleep(2)
|
|
print(f" [TIMEOUT] {prompt_id}")
|
|
return False
|
|
|
|
|
|
def free_memory():
|
|
try:
|
|
req = urllib.request.Request(
|
|
f"{COMFY_URL}/free",
|
|
data=json.dumps({"unload_models": True, "free_memory": True}).encode("utf-8"),
|
|
headers={"Content-Type": "application/json"},
|
|
)
|
|
urllib.request.urlopen(req)
|
|
time.sleep(3)
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
# ── File Discovery ───────────────────────────────────────────
|
|
|
|
def find_image(directory, file_tag, prompt_id, seed):
|
|
exact = os.path.join(directory, f"{prompt_id}_{file_tag}_seed{seed}_00001_.png")
|
|
if os.path.exists(exact):
|
|
return exact
|
|
pattern = os.path.join(directory, f"{prompt_id}_{file_tag}_seed{seed}_*.png")
|
|
matches = glob.glob(pattern)
|
|
return matches[0] if matches else None
|
|
|
|
|
|
def image_exists(directory, file_tag, prompt_id, seed):
|
|
return find_image(directory, file_tag, prompt_id, seed) is not None
|
|
|
|
|
|
def count_images(directory, file_tag, prompt_ids, seeds):
|
|
return sum(1 for p in prompt_ids for s in seeds
|
|
if find_image(directory, file_tag, p, s))
|
|
|
|
|
|
# ── Resume Support ───────────────────────────────────────────
|
|
|
|
class ProgressLog:
|
|
def __init__(self, path):
|
|
self.path = path
|
|
self.done = set()
|
|
if os.path.exists(path):
|
|
with open(path, "r") as f:
|
|
self.done = {line.strip() for line in f if line.strip()}
|
|
|
|
def is_done(self, key):
|
|
return key in self.done
|
|
|
|
def mark_done(self, key):
|
|
self.done.add(key)
|
|
with open(self.path, "a") as f:
|
|
f.write(key + "\n")
|
|
|
|
|
|
# ── Workflow Builders ────────────────────────────────────────
|
|
|
|
def build_uniform_workflow(model_file, prompt_text, seed, filename_prefix,
|
|
steps=20, cfg=1.0, guidance=3.5,
|
|
t5=None, clip_l=None, vae=None):
|
|
from config import T5_ENCODER, CLIP_L, VAE, SAMPLER, SCHEDULER, WIDTH, HEIGHT
|
|
t5 = t5 or T5_ENCODER
|
|
clip_l = clip_l or CLIP_L
|
|
vae = vae or VAE
|
|
return {
|
|
"1": {"class_type": "DualCLIPLoaderGGUF",
|
|
"inputs": {"clip_name1": t5, "clip_name2": clip_l, "type": "flux"}},
|
|
"2": {"class_type": "UnetLoaderGGUF",
|
|
"inputs": {"unet_name": model_file}},
|
|
"3": {"class_type": "VAELoader",
|
|
"inputs": {"vae_name": vae}},
|
|
"4": {"class_type": "VAEDecode",
|
|
"inputs": {"samples": ["7", 0], "vae": ["3", 0]}},
|
|
"5": {"class_type": "CLIPTextEncodeFlux",
|
|
"inputs": {"clip_l": prompt_text, "t5xxl": prompt_text,
|
|
"guidance": guidance, "clip": ["1", 0]}},
|
|
"6": {"class_type": "EmptyLatentImage",
|
|
"inputs": {"width": WIDTH, "height": HEIGHT, "batch_size": 1}},
|
|
"7": {"class_type": "KSampler",
|
|
"inputs": {"seed": seed, "steps": steps, "cfg": cfg,
|
|
"sampler_name": SAMPLER, "scheduler": SCHEDULER,
|
|
"denoise": 1, "model": ["2", 0],
|
|
"positive": ["5", 0], "negative": ["5", 0],
|
|
"latent_image": ["6", 0]}},
|
|
"8": {"class_type": "SaveImage",
|
|
"inputs": {"filename_prefix": filename_prefix, "images": ["4", 0]}},
|
|
}
|
|
|
|
|
|
def build_dual_workflow(model_a, model_b, switch_step, prompt_text, seed,
|
|
filename_prefix, steps=20, cfg=1.0, guidance=3.5,
|
|
t5=None, clip_l=None, vae=None):
|
|
from config import T5_ENCODER, CLIP_L, VAE, SAMPLER, SCHEDULER, WIDTH, HEIGHT
|
|
t5 = t5 or T5_ENCODER
|
|
clip_l = clip_l or CLIP_L
|
|
vae = vae or VAE
|
|
return {
|
|
"1": {"class_type": "DualModelKSampler",
|
|
"inputs": {"seed": seed, "steps": steps, "cfg": cfg,
|
|
"sampler_name": SAMPLER, "scheduler": SCHEDULER,
|
|
"switch_step": switch_step, "denoise": 1, "switch_end": 0,
|
|
"model_a": ["3", 0], "model_b": ["4", 0],
|
|
"positive": ["7", 0], "negative": ["7", 0],
|
|
"latent_image": ["8", 0]}},
|
|
"3": {"class_type": "UnetLoaderGGUF", "inputs": {"unet_name": model_a}},
|
|
"4": {"class_type": "UnetLoaderGGUF", "inputs": {"unet_name": model_b}},
|
|
"6": {"class_type": "VAELoader", "inputs": {"vae_name": vae}},
|
|
"7": {"class_type": "CLIPTextEncodeFlux",
|
|
"inputs": {"clip_l": prompt_text, "t5xxl": prompt_text,
|
|
"guidance": guidance, "clip": ["12", 0]}},
|
|
"8": {"class_type": "EmptyLatentImage",
|
|
"inputs": {"width": WIDTH, "height": HEIGHT, "batch_size": 1}},
|
|
"9": {"class_type": "VAEDecode",
|
|
"inputs": {"samples": ["1", 0], "vae": ["6", 0]}},
|
|
"10": {"class_type": "SaveImage",
|
|
"inputs": {"filename_prefix": filename_prefix, "images": ["9", 0]}},
|
|
"12": {"class_type": "DualCLIPLoaderGGUF",
|
|
"inputs": {"clip_name1": t5, "clip_name2": clip_l, "type": "flux"}},
|
|
}
|
|
|
|
|
|
def build_stepwise_workflow(model_file, prompt_text, seed, save_prefix,
|
|
steps=20, cfg=1.0, guidance=3.5,
|
|
save_steps="all",
|
|
t5=None, clip_l=None, vae=None):
|
|
from config import T5_ENCODER, CLIP_L, VAE, SAMPLER, SCHEDULER, WIDTH, HEIGHT
|
|
t5 = t5 or T5_ENCODER
|
|
clip_l = clip_l or CLIP_L
|
|
vae = vae or VAE
|
|
return {
|
|
"1": {"class_type": "DualCLIPLoaderGGUF",
|
|
"inputs": {"clip_name1": t5, "clip_name2": clip_l, "type": "flux"}},
|
|
"2": {"class_type": "UnetLoaderGGUF",
|
|
"inputs": {"unet_name": model_file}},
|
|
"3": {"class_type": "VAELoader",
|
|
"inputs": {"vae_name": vae}},
|
|
"4": {"class_type": "VAEDecode",
|
|
"inputs": {"samples": ["7", 0], "vae": ["3", 0]}},
|
|
"5": {"class_type": "CLIPTextEncodeFlux",
|
|
"inputs": {"clip_l": prompt_text, "t5xxl": prompt_text,
|
|
"guidance": guidance, "clip": ["1", 0]}},
|
|
"6": {"class_type": "EmptyLatentImage",
|
|
"inputs": {"width": WIDTH, "height": HEIGHT, "batch_size": 1}},
|
|
"7": {"class_type": "StepwiseKSampler",
|
|
"inputs": {"model": ["2", 0], "positive": ["5", 0],
|
|
"negative": ["5", 0], "latent_image": ["6", 0],
|
|
"seed": seed, "steps": steps, "cfg": cfg,
|
|
"sampler_name": SAMPLER, "scheduler": SCHEDULER,
|
|
"denoise": 1.0, "save_steps": save_steps,
|
|
"save_prefix": save_prefix}},
|
|
"8": {"class_type": "SaveImage",
|
|
"inputs": {"filename_prefix": save_prefix, "images": ["4", 0]}},
|
|
}
|
|
|
|
|
|
# ── Stats Helpers ────────────────────────────────────────────
|
|
|
|
def avg(lst):
|
|
return sum(lst) / len(lst) if lst else 0.0
|
|
|
|
def std(lst):
|
|
if not lst:
|
|
return 0.0
|
|
m = avg(lst)
|
|
return (sum((v - m) ** 2 for v in lst) / len(lst)) ** 0.5
|