92 lines
3.5 KiB
Python
92 lines
3.5 KiB
Python
"""Encola N variantes de un workflow ComfyUI, una por seed, y recoge los prompt_ids.
|
|
|
|
Funcion impura: hace red (POST /prompt por variante, via comfyui_submit_workflow).
|
|
Compone comfyui_submit_workflow.
|
|
|
|
Para cada seed de la lista, copia el workflow (deepcopy, no muta el original),
|
|
parchea el campo de semilla de los nodos sampler (KSampler.seed, KSamplerAdvanced.
|
|
noise_seed, SamplerCustom.noise_seed — en general cualquier input "seed"/"noise_seed")
|
|
y lo encola. Util para barridos de re-roll: misma escena, varias semillas, una sola
|
|
llamada. Devuelve los prompt_ids en el mismo orden que la lista de seeds; cada uno
|
|
se sigue con comfyui_wait_result.
|
|
"""
|
|
import copy
|
|
import os
|
|
import sys
|
|
|
|
_THIS_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
if _THIS_DIR not in sys.path:
|
|
sys.path.insert(0, _THIS_DIR)
|
|
|
|
from comfyui_submit_workflow import comfyui_submit_workflow # noqa: E402
|
|
|
|
# Campos de semilla conocidos en los nodos sampler de ComfyUI.
|
|
_SEED_KEYS = ("seed", "noise_seed")
|
|
|
|
|
|
def _patch_seed(workflow: dict, seed: int) -> dict:
|
|
"""Copia el workflow y fija `seed` en todos los inputs de semilla (no muta el original)."""
|
|
wf = copy.deepcopy(workflow)
|
|
for node in wf.values():
|
|
inputs = node.get("inputs")
|
|
if not isinstance(inputs, dict):
|
|
continue
|
|
for key in _SEED_KEYS:
|
|
if key in inputs:
|
|
inputs[key] = seed
|
|
return wf
|
|
|
|
|
|
def comfyui_batch_generate(
|
|
workflow: dict,
|
|
*,
|
|
seeds: list | None = None,
|
|
server: str = "127.0.0.1:8188",
|
|
) -> dict:
|
|
"""Encola una variante del workflow por cada seed y devuelve los prompt_ids.
|
|
|
|
Args:
|
|
workflow: dict en API format (resultado de un builder). No se muta: cada
|
|
variante es una copia profunda con la semilla parcheada.
|
|
seeds: lista de semillas (int). Cada una produce una variante encolada. Si
|
|
es None o vacia, se encola el workflow tal cual una sola vez (sin
|
|
parchear semilla). keyword-only.
|
|
server: host:port del servidor ComfyUI sin esquema. keyword-only.
|
|
|
|
Returns:
|
|
dict con:
|
|
- ok (bool): True si TODAS las variantes se encolaron sin error.
|
|
- prompt_ids (list[str]): prompt_id de cada variante encolada, en orden.
|
|
- count (int): numero de variantes encoladas con exito.
|
|
- error (str): primer error encontrado; cadena vacia si todo OK. Si una
|
|
variante falla, se detiene el barrido y se devuelven los prompt_ids ya
|
|
encolados.
|
|
"""
|
|
out = {"ok": False, "prompt_ids": [], "count": 0, "error": ""}
|
|
variants = [(s, _patch_seed(workflow, s)) for s in seeds] if seeds else [(None, workflow)]
|
|
|
|
for seed, wf in variants:
|
|
try:
|
|
resp = comfyui_submit_workflow(wf, server=server)
|
|
except RuntimeError as exc:
|
|
label = "tal cual" if seed is None else f"seed={seed}"
|
|
out["error"] = f"variante {label} fallo al encolar: {exc}"
|
|
return out
|
|
out["prompt_ids"].append(resp["prompt_id"])
|
|
|
|
out["count"] = len(out["prompt_ids"])
|
|
out["ok"] = True
|
|
return out
|
|
|
|
|
|
if __name__ == "__main__":
|
|
from comfyui_build_txt2img_workflow import comfyui_build_txt2img_workflow
|
|
|
|
wf = comfyui_build_txt2img_workflow(
|
|
ckpt_name="v1-5-pruned-emaonly-fp16.safetensors",
|
|
positive="a red apple on a wooden table, sharp focus",
|
|
negative="blurry, low quality",
|
|
)
|
|
res = comfyui_batch_generate(wf, seeds=[1, 2])
|
|
print(f"ok={res['ok']} count={res['count']} ids={res['prompt_ids']} error={res['error']!r}")
|