75 lines
2.6 KiB
Python
75 lines
2.6 KiB
Python
"""
|
||
Generate uniform quantization ladder images.
|
||
8 quant levels × 25 prompts × 10 seeds = 2,000 images.
|
||
"""
|
||
|
||
import argparse
|
||
import os
|
||
import time
|
||
from config import GGUF_MODELS, MODEL_ORDER, PROMPTS, SEEDS
|
||
from utils import check_comfyui, queue_prompt, wait_for_completion, \
|
||
build_uniform_workflow, image_exists, ProgressLog
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(description="Generate uniform quant ladder")
|
||
parser.add_argument("--output-dir", required=True)
|
||
parser.add_argument("--comfyui-output-dir", default=None,
|
||
help="ComfyUI's actual output directory (for resume check)")
|
||
args = parser.parse_args()
|
||
|
||
os.makedirs(args.output_dir, exist_ok=True)
|
||
check_dir = args.comfyui_output_dir or args.output_dir
|
||
check_comfyui()
|
||
|
||
log = ProgressLog(os.path.join(args.output_dir, "completed.log"))
|
||
total = len(MODEL_ORDER) * len(PROMPTS) * len(SEEDS)
|
||
done, skipped = 0, 0
|
||
start = time.time()
|
||
|
||
print(f"Total: {total} images ({len(MODEL_ORDER)} models × "
|
||
f"{len(PROMPTS)} prompts × {len(SEEDS)} seeds)\n")
|
||
|
||
for model_tag in MODEL_ORDER:
|
||
model_file = GGUF_MODELS[model_tag]
|
||
print(f"\n{'='*50}\n Model: {model_tag} ({model_file})\n{'='*50}")
|
||
|
||
for pid, ptext in PROMPTS.items():
|
||
for seed in SEEDS:
|
||
key = f"{pid}_{model_tag}_seed{seed}"
|
||
done += 1
|
||
|
||
if log.is_done(key):
|
||
skipped += 1
|
||
continue
|
||
|
||
elapsed = time.time() - start
|
||
generated = done - skipped
|
||
eta = (total - done) * (elapsed / max(generated, 1))
|
||
print(f" [{done}/{total}] {key} "
|
||
f"elapsed={elapsed/3600:.1f}h eta={eta/3600:.1f}h",
|
||
end=" ", flush=True)
|
||
|
||
wf = build_uniform_workflow(model_file, ptext, seed, key)
|
||
t0 = time.time()
|
||
try:
|
||
prompt_id = queue_prompt(wf)
|
||
ok = wait_for_completion(prompt_id)
|
||
dt = time.time() - t0
|
||
if ok:
|
||
log.mark_done(key)
|
||
print(f"ok ({dt:.1f}s)")
|
||
else:
|
||
print(f"FAIL ({dt:.1f}s)")
|
||
except Exception as e:
|
||
print(f"ERROR: {e}")
|
||
time.sleep(5)
|
||
|
||
total_time = time.time() - start
|
||
generated = done - skipped
|
||
print(f"\nDone: {generated} generated, {skipped} skipped, {total_time/3600:.1f}h total")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|