test_pipelines.py

"""
End-to-end test: run every output in workflow.json through the real
`WorkflowExecutor` — the same code path the canvas and the REST API use.

    python apps/05_workflow1111/test_pipelines.py             # everything
    python apps/05_workflow1111/test_pipelines.py local       # only offline nodes
    python apps/05_workflow1111/test_pipelines.py grid image  # substring filters

Hits Hugging Face for the `model`/`space` nodes, so it costs quota and takes a
couple of minutes. Rendered outputs are written to ./_test_output for eyeballing.
"""

import base64
import inspect
import json
import logging
import os
import sys
import time
import types
import urllib.request
import warnings

warnings.filterwarnings("ignore")
logging.disable(logging.CRITICAL)

for _s in (sys.stdout, sys.stderr):
    try:
        _s.reconfigure(encoding="utf-8", errors="replace")
    except (AttributeError, ValueError):
        pass

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)

import gradio.workflow as W  # noqa: E402
from gradio.helpers import special_args  # noqa: E402
from gradio.workflow_api import (  # noqa: E402
    WorkflowExecutor,
    WorkflowGraph,
    group_free_inputs,
    subject_groups,
)
from huggingface_hub import get_token  # noqa: E402

import nodes as N  # noqa: E402

TOKEN = types.SimpleNamespace(token=get_token() or os.environ.get("HF_TOKEN"))
OUTDIR = os.path.join(HERE, "_test_output")
os.makedirs(OUTDIR, exist_ok=True)

# A real photograph — DETR/ViT/the VLM need actual content to say anything about.
SAMPLE_URL = "https://huggingface.co/datasets/mishig/sample_images/resolve/main/tiger.jpg"
SAMPLE = os.path.join(OUTDIR, "_sample.jpg")
if not os.path.exists(SAMPLE):
    urllib.request.urlretrieve(SAMPLE_URL, SAMPLE)  # noqa: S310

# An image that carries embedded generation parameters, for the PNG Info pipeline.
STAMPED = os.path.join(OUTDIR, "_stamped.png")
if not os.path.exists(STAMPED):
    _info = N.generation_info("a red fox in snow", "blurry, watermark", 8, 3.5,
                              987654, 1024, 768, "black-forest-labs/FLUX.1-schnell")
    _uri = N.postprocess(N._emit(N._load_image(SAMPLE)), 1, "Lanczos", 0, 1, 1, 1,
                         0, 0, 0, "", _info)
    with open(STAMPED, "wb") as f:
        f.write(base64.b64decode(_uri.partition(",")[2]))


def call_fn(data, request=None, token=None):
    """Mirror of gradio's bound-function server fn (workflow.py `call_fn`)."""
    name = data[0] if data else ""
    fn = N.BIND.get(name)
    if fn is None:
        return json.dumps({"error": f"No function '{name}' bound"})
    try:
        args = json.loads(data[1] if len(data) > 1 else "[]")
        if not isinstance(args, list):
            args = [args]
        # gradio injects OAuthToken/Request params before calling — mirror that
        # so `txt2img` receives its token exactly as it will in the app.
        args, *_ = special_args(fn, args, request, None, token=token)
        result = fn(*args)
        return json.dumps(list(result) if isinstance(result, (list, tuple)) else [result])
    except Exception as e:
        return json.dumps({"error": f"{type(e).__name__}: {e}"})


CALLERS = {"fn": call_fn, "model": W.call_model,
           "space": W.call_space, "dataset": W.fetch_dataset}

with open(os.path.join(HERE, "workflow.json"), encoding="utf-8") as f:
    GRAPH = WorkflowGraph.from_json(f.read())

# Values fed to reference nodes. Anything not listed falls back to the node's
# own default (`data.out`), or to the sample image for media ports.
OVERRIDES = {
    "ref_interrogate_image": SAMPLE,
    "ref_detect_image": SAMPLE,
    "ref_extras_image": SAMPLE,
    "ref_control_image": SAMPLE,
    "ref_init_image": SAMPLE,
    "ref_pnginfo_image": STAMPED,
}

# Subjects whose upstream is entirely local — these must pass with no network.
LOCAL_ONLY = {"sub_control", "sub_png_report", "sub_png_fields",
              "sub_upscaled", "sub_upscale_report"}


def seed_for(subject_id):
    node = GRAPH.node_by_id[subject_id]
    inputs = {}
    # free_inputs yields {"node", "port", "type", "label"} wrappers, not the
    # reference nodes themselves.
    for free in group_free_inputs(GRAPH, [node]):
        ref = free["node"]
        rid = ref["id"]
        if rid in OVERRIDES:
            inputs[rid] = OVERRIDES[rid]
            continue
        default = (ref.get("data") or {}).get("out")
        if default in (None, "") and free["type"] in ("image", "audio", "video"):
            default = SAMPLE
        inputs[rid] = default
    return inputs


def save(subject_id, value):
    """Persist an output so it can actually be looked at; return a summary."""
    if isinstance(value, str) and value.startswith("data:"):
        header, _, payload = value.partition(",")
        ext = "png" if "png" in header else "jpg"
        path = os.path.join(OUTDIR, f"{subject_id}.{ext}")
        raw = base64.b64decode(payload)
        with open(path, "wb") as f:
            f.write(raw)
        from PIL import Image
        with Image.open(path) as im:
            return f"image {im.width}×{im.height} {ext.upper()}, {len(raw) // 1024} KB"

    if isinstance(value, str) and os.path.isfile(value):
        from PIL import Image
        try:
            with Image.open(value) as im:
                dst = os.path.join(OUTDIR, f"{subject_id}.png")
                im.convert("RGBA" if im.mode in ("RGBA", "LA", "P") else "RGB").save(dst)
                return f"image {im.width}×{im.height} (file)"
        except Exception:
            return f"file {os.path.basename(value)}"

    if isinstance(value, (dict, list)):
        path = os.path.join(OUTDIR, f"{subject_id}.json")
        with open(path, "w", encoding="utf-8") as f:
            json.dump(value, f, indent=2, ensure_ascii=False)
        return f"{type(value).__name__} ({len(value)} entries)"

    text = str(value)
    path = os.path.join(OUTDIR, f"{subject_id}.txt")
    with open(path, "w", encoding="utf-8") as f:
        f.write(text)
    return f"text ({len(text)} chars): {text.splitlines()[0][:76] if text.strip() else '(empty)'}"


def main():
    args = [a.lower() for a in sys.argv[1:]]
    local_only = "local" in args
    filters = [a for a in args if a != "local"]

    executor = WorkflowExecutor(GRAPH, CALLERS)
    targets = [s["id"] for s in GRAPH.subjects]
    if local_only:
        targets = [t for t in targets if t in LOCAL_ONLY]
    if filters:
        targets = [t for t in targets
                   if any(f in t.lower() or f in GRAPH.node_by_id[t]["label"].lower()
                          for f in filters)]

    if not targets:
        print("no subjects matched")
        return 1

    print(f"\nRunning {len(targets)} output(s) through WorkflowExecutor")
    print(f"outputs → {os.path.relpath(OUTDIR, os.getcwd())}\n")

    passed, failed = [], []
    for sid in targets:
        label = GRAPH.node_by_id[sid]["label"]
        t0 = time.time()
        try:
            value = executor.run(sid, seed_for(sid), request=None, token=TOKEN)
            if value is None or value == "":
                raise AssertionError("output was empty")
            summary = save(sid, value)
        except Exception as e:
            failed.append((sid, f"{type(e).__name__}: {e}"))
            print(f"  FAIL  {label:28} ({time.time() - t0:6.1f}s)  "
                  f"{type(e).__name__}: {str(e)[:130]}")
        else:
            passed.append(sid)
            print(f"  ok    {label:28} ({time.time() - t0:6.1f}s)  {summary}")
        sys.stdout.flush()

    print("\n── API surface (one endpoint per subject group) ──")
    for group in subject_groups(GRAPH):
        names = ", ".join(s["label"] for s in group)
        params = ", ".join(f"{f['label']} ({f['type']})"
                           for f in group_free_inputs(GRAPH, group))
        print(f"  • {names}\n      inputs: {params or '(none)'}")

    print(f"\n{'=' * 66}\n  {len(passed)} passed, {len(failed)} failed\n{'=' * 66}")
    for sid, err in failed:
        print(f"  {sid}: {err}")
    return 1 if failed else 0


if __name__ == "__main__":
    sys.exit(main())