fn_registry/python/functions/pipelines/comfyui_txt2img_oneshot.py

"""comfyui_txt2img_oneshot — prompt de texto -> PNG en disco en una sola llamada.

Promocion de la secuencia repetida (issue 0087): construir el workflow txt2img ->
encolar -> esperar -> descargar la imagen. Compone funciones del registry del
grupo `comfyui`:

    comfyui_build_txt2img_workflow_py_ml  (workflow de nodos en API format)
    comfyui_submit_workflow_py_ml         (POST /prompt)
    comfyui_wait_result_py_ml             (poll /history)
    comfyui_fetch_output_image_py_ml      (GET /view -> disco)

Pipeline impuro: red (HTTP) + escritura en disco.
"""

from __future__ import annotations

import os
import sys

# Importa las funciones del registry (mismo arbol python/functions).
_FUNCTIONS_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if _FUNCTIONS_ROOT not in sys.path:
    sys.path.insert(0, _FUNCTIONS_ROOT)

from ml.comfyui_build_txt2img_workflow import comfyui_build_txt2img_workflow
from ml.comfyui_fetch_output_image import comfyui_fetch_output_image
from ml.comfyui_submit_workflow import comfyui_submit_workflow
from ml.comfyui_wait_result import comfyui_wait_result


def comfyui_txt2img_oneshot(
    prompt: str,
    *,
    ckpt: str = "IMG_dreamshaper_8.safetensors",
    negative: str = "",
    server: str = "127.0.0.1:8188",
    dest: str | None = None,
    wait_timeout: float = 300.0,
    **gen,
) -> dict:
    """Genera una imagen desde un prompt de texto, end-to-end.

    Args:
        prompt: prompt positivo (lo que se quiere ver en la imagen).
        ckpt: checkpoint Stable Diffusion tal como lo ve el servidor
            (CheckpointLoaderSimple). Por defecto "IMG_dreamshaper_8.safetensors".
            keyword-only.
        negative: prompt negativo. Por defecto "". keyword-only.
        server: host:port del servidor ComfyUI (sin esquema). keyword-only.
        dest: directorio local donde guardar el PNG (None = cwd). keyword-only.
        wait_timeout: segundos maximos esperando a que el server termine.
            keyword-only.
        **gen: parametros de generacion pasados a comfyui_build_txt2img_workflow
            (steps, cfg, width, height, seed, sampler_name, scheduler,
            filename_prefix).

    Returns:
        dict {ok, image_path, prompt_id, error}. image_path = ruta local del PNG
        descargado; prompt_id = id del trabajo en ComfyUI. Si falla, ok=False y
        error explica en que paso.
    """
    # 1. Construir el workflow (funcion pura del registry).
    workflow = comfyui_build_txt2img_workflow(ckpt, prompt, negative, **gen)

    # 2. Encolar.
    try:
        sub = comfyui_submit_workflow(workflow, server=server)
        prompt_id = sub["prompt_id"]
    except (RuntimeError, KeyError) as exc:
        return {"ok": False, "image_path": "", "prompt_id": "",
                "error": f"submit fallo: {exc}"}

    # 3. Esperar a que termine.
    try:
        outputs = comfyui_wait_result(prompt_id, server=server, timeout=wait_timeout)
    except (TimeoutError, RuntimeError) as exc:
        return {"ok": False, "image_path": "", "prompt_id": prompt_id,
                "error": f"wait fallo: {exc}"}

    # 4. Localizar el primer PNG en los outputs (nodo SaveImage -> images).
    img = None
    for node_out in outputs.values():
        images = node_out.get("images") if isinstance(node_out, dict) else None
        if images:
            img = images[0]
            break
    if img is None:
        return {"ok": False, "image_path": "", "prompt_id": prompt_id,
                "error": f"el workflow no produjo imagenes (outputs={list(outputs)})"}

    # 5. Descargar la imagen a disco.
    fetched = comfyui_fetch_output_image(
        img["filename"],
        subfolder=img.get("subfolder", ""),
        type_=img.get("type", "output"),
        server=server,
        dest_dir=dest or ".",
    )
    if not fetched.get("ok"):
        return {"ok": False, "image_path": "", "prompt_id": prompt_id,
                "error": f"fetch de imagen fallo: {fetched.get('error')}"}

    return {"ok": True, "image_path": fetched["path"], "prompt_id": prompt_id,
            "error": ""}


if __name__ == "__main__":
    import json

    res = comfyui_txt2img_oneshot(
        "a red apple on a wooden table, sharp focus",
        negative="blurry, low quality",
        dest="/tmp/comfy_txt2img",
        steps=20,
        seed=42,
    )
    print(json.dumps(res, indent=2))