diff --git a/README.md b/README.md index 2f047e6b..a27c72f1 100644 --- a/README.md +++ b/README.md @@ -19,7 +19,7 @@ workflows, and call hosted partner image models, all from your terminal. ## Features - πŸš€ One-command ComfyUI install and launch -- 🎨 Direct calls to partner image nodes (Flux, Ideogram, DALLΒ·E, Recraft, Stability, …) via `comfy generate`, no workflow JSON required +- 🎨 Direct calls to partner image and video nodes (Flux, Ideogram, DALLΒ·E, Recraft, Stability, Kling, Luma, Runway, Pika, Vidu, Hailuo, …) via `comfy generate`, no workflow JSON required - πŸ”§ Custom node management β€” install, update, snapshot, bisect - πŸ“¦ Fast dependency resolution with `uv` (`--fast-deps`, `--uv-compile`) - πŸ—„οΈ Model downloads from CivitAI, Hugging Face, and direct URLs @@ -290,8 +290,11 @@ the bisect tool can help you pinpoint the custom node that causes the issue. `comfy generate` calls Comfy's partner nodes directly from the terminal β€” no local ComfyUI or workflow JSON required. It hits the same hosted partner nodes -(Flux, Ideogram, DALLΒ·E, Recraft, Stability, Runway, Reve, xAI Grok, …) you'd -otherwise wire into a ComfyUI workflow, but as one-shot CLI calls. +you'd otherwise wire into a ComfyUI workflow, but as one-shot CLI calls. Image +models (Flux, Ideogram, DALLΒ·E, Recraft, Stability, Runway, Reve, xAI Grok, …) +and video models (Kling, Luma, Runway Gen-3, Pika, Vidu, Moonvalley, Hailuo, +Grok video) are all covered; video jobs run async and the CLI polls until the +result is ready. Prerequisites β€” a Comfy API key and a credit balance: @@ -321,8 +324,18 @@ comfy generate flux-kontext --prompt "add a top hat" \ comfy generate upload ./photo.jpg # explicit upload ``` -Async models block until ready by default. Pass `--async` to return immediately -with a job id, then resume later with `comfy generate resume `. +Async models (every video model plus the Flux family) block until ready by +default. Pass `--async` to return immediately with a job id, then resume later +with `comfy generate resume `. Examples: + +```bash +comfy generate kling --prompt "a paper boat drifting on a river at dusk" \ + --duration 5 --download boat.mp4 + +comfy generate luma --prompt "..." --aspect_ratio 16:9 --async +# β†’ prints job id; resume with: +comfy generate resume luma --download out.mp4 +``` ### Managing ComfyUI-Manager diff --git a/comfy_cli/command/generate/app.py b/comfy_cli/command/generate/app.py index 1897411a..5897c009 100644 --- a/comfy_cli/command/generate/app.py +++ b/comfy_cli/command/generate/app.py @@ -223,7 +223,7 @@ def _generate(model: str, extra_args: list[str]) -> None: raise typer.Exit(code=1) if ep.polling: - job_id = str(body.get("id") or (body.get("data") or {}).get("task_id") or request_id) + job_id = poll.extract_job_id(ep.polling, body) or request_id name = spec.preferred_alias(ep.id) or ep.id if do_async: if as_json: @@ -241,7 +241,13 @@ def _generate(model: str, extra_args: list[str]) -> None: def _on_progress(p: float) -> None: prog.update(task, description=f"Generating ({p * 100:.0f}%)") - result = poller(body, api_key=api_key, timeout=timeout, on_progress=_on_progress) + result = poller( + body, + api_key=api_key, + timeout=timeout, + on_progress=_on_progress, + create_path=ep.path, + ) _emit_result(result, request_id=job_id, download=download, as_json=as_json) return @@ -411,10 +417,10 @@ def _resume(extra_args: list[str]) -> None: download = meta.get("download") if isinstance(meta.get("download"), str) else None as_json = bool(meta.get("json", False)) - if ep.polling == "bfl": - initial = {"polling_url": f"{spec.base_url()}/proxy/bfl/get_result?id={job_id}"} - else: - rprint(f"[bold red]Resume not implemented for partner {ep.partner}[/bold red]") + try: + initial = poll.build_synthetic_initial(ep.polling, job_id, base_url=spec.base_url()) + except client.ApiError as e: + rprint(f"[bold red]{e}[/bold red]") raise typer.Exit(code=1) poller = poll.get_poller(ep.polling) @@ -424,7 +430,13 @@ def _resume(extra_args: list[str]) -> None: def _on_progress(p: float) -> None: prog.update(task, description=f"Job {job_id} ({p * 100:.0f}%)") - result = poller(initial, api_key=api_key, timeout=timeout, on_progress=_on_progress) + result = poller( + initial, + api_key=api_key, + timeout=timeout, + on_progress=_on_progress, + create_path=ep.path, + ) _emit_result(result, request_id=job_id, download=download, as_json=as_json) diff --git a/comfy_cli/command/generate/output.py b/comfy_cli/command/generate/output.py index 1a1f042d..da0ecd8b 100644 --- a/comfy_cli/command/generate/output.py +++ b/comfy_cli/command/generate/output.py @@ -48,11 +48,21 @@ def _resolve_template(template: str, request_id: str, index: int, ext: str) -> P def save_urls(urls: list[str], template: str, request_id: str) -> list[Path]: - """Download each URL and save under the resolved template path. Returns saved paths.""" + """Download each URL and save under the resolved template path. Returns saved paths. + + Multi-URL responses (a video + thumbnail from Luma, for example) need a + per-output filename. If the template has no ``{index}`` placeholder and + isn't a directory shorthand, we auto-insert ``_`` before the suffix + and switch the extension to whatever the URL actually points at β€” so a + user who typed ``--download out.mp4`` doesn't silently get a thumbnail + JPEG written into ``out.mp4`` because the model returned two URLs.""" saved: list[Path] = [] + auto_index = len(urls) > 1 and "{index}" not in template and not template.endswith(("/", "\\")) for i, url in enumerate(urls): ext = _ext_from_url(url) dest = _resolve_template(template, request_id, i, ext) + if auto_index: + dest = dest.with_name(f"{dest.stem}_{i}.{ext}") dest.parent.mkdir(parents=True, exist_ok=True) data = client.download_bytes(url) dest.write_bytes(data) diff --git a/comfy_cli/command/generate/poll.py b/comfy_cli/command/generate/poll.py index 2cd64ccf..8f4edbee 100644 --- a/comfy_cli/command/generate/poll.py +++ b/comfy_cli/command/generate/poll.py @@ -1,14 +1,24 @@ -"""Async-job polling adapters, one per partner. +"""Async-job polling for partner endpoints. -Each adapter takes the initial response from the submit POST and yields a -canonical ``PollResult`` once the upstream job reaches a terminal state. +There are two flavors: + +1. **BFL** β€” the server returns ``{id, polling_url}`` on submit and we just GET + that URL until the ``status`` field is terminal. +2. **Everything else** β€” a small ``PollSpec`` per partner describes where the + job id lives in the create response, how to construct the poll URL (some + partners use a sibling endpoint relative to the create path; others have a + dedicated ``/tasks/{id}`` endpoint), and which status values mean + "succeeded" / "failed". + +The generic poller walks dot-paths into the JSON to extract the id/status +without having to write a new adapter for each partner. """ from __future__ import annotations import time from collections.abc import Callable -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import Any import httpx @@ -22,23 +32,46 @@ class PollResult: status: str # "succeeded" | "failed" | "cancelled" raw: dict[str, Any] # last response body β€” full upstream payload - image_urls: list[str] # any image result URLs we could pluck out + image_urls: list[str] # any image/video result URLs we could pluck out error: str | None = None +# Recognized result extensions when sniffing URLs out of a poll body. +_MEDIA_EXTS = ( + ".png", + ".jpg", + ".jpeg", + ".webp", + ".gif", + ".svg", + ".mp4", + ".mov", + ".webm", + ".m4v", + ".gltf", + ".glb", + ".obj", + ".fbx", + ".wav", + ".mp3", + ".m4a", + ".flac", +) + + def _now() -> float: return time.monotonic() def _extract_urls(node: Any) -> list[str]: - """Walk a JSON tree, collecting strings that look like image URLs.""" + """Walk a JSON tree, collecting strings that look like media URLs.""" found: list[str] = [] def visit(n: Any) -> None: if isinstance(n, str): low = n.lower() if n.startswith(("http://", "https://")) and ( - low.endswith((".png", ".jpg", ".jpeg", ".webp", ".gif", ".svg")) or "image" in low + low.split("?", 1)[0].endswith(_MEDIA_EXTS) or "image" in low or "video" in low ): found.append(n) return @@ -50,11 +83,29 @@ def visit(n: Any) -> None: visit(v) visit(node) - # De-dupe preserving order. seen: set[str] = set() return [u for u in found if not (u in seen or seen.add(u))] +def _dotget(body: Any, path: str) -> Any: + """Look up a dotted path inside a JSON body. Returns None if any segment misses.""" + cur: Any = body + for part in path.split("."): + if isinstance(cur, dict): + cur = cur.get(part) + else: + return None + return cur + + +def _first(body: Any, paths: tuple[str, ...]) -> Any: + for p in paths: + v = _dotget(body, p) + if v not in (None, "", []): + return v + return None + + def _sleep(seconds: float) -> None: time.sleep(seconds) @@ -66,8 +117,9 @@ def poll_bfl( interval: float = 2.0, timeout: float = 300.0, on_progress: Callable[[float], None] | None = None, + create_path: str | None = None, # ignored, kept for uniform signature ) -> PollResult: - """BFL: GET ``polling_url`` until ``status`` is terminal.""" + """BFL polls a server-issued ``polling_url`` until ``status`` flips to Ready.""" url = initial.get("polling_url") if not url: raise client.ApiError(0, "", "BFL response missing polling_url") @@ -92,24 +144,235 @@ def poll_bfl( return PollResult(status="failed", raw=last_body, image_urls=[], error=f"timed out after {timeout:.0f}s") -_POLLERS: dict[str, Callable[..., PollResult]] = { - "bfl": poll_bfl, - # "kling": poll_kling, # follow-up: not in v1 image allowlist - # "luma": poll_luma, # follow-up - # "topaz": poll_topaz, # follow-up +@dataclass(frozen=True) +class PollSpec: + """Per-partner polling configuration. + + ``poll_url`` is a template supporting ``{id}`` and ``{create_path}``; + ``post_success_url`` (optional) is a second-stage fetcher invoked once the + job reaches a success state β€” for partners like MiniMax where the terminal + poll response gives you a file id you still need to redeem for a URL.""" + + name: str + id_paths: tuple[str, ...] + poll_url: str + status_paths: tuple[str, ...] + success_values: tuple[str, ...] + failure_values: tuple[str, ...] = () + progress_path: str | None = None + post_success_url: str | None = None + post_success_id_paths: tuple[str, ...] = field(default_factory=tuple) + + +_POLL_SPECS: dict[str, PollSpec] = { + "kling": PollSpec( + name="kling", + id_paths=("data.task_id",), + poll_url="{create_path}/{id}", + status_paths=("data.task_status",), + success_values=("succeed",), + failure_values=("failed",), + ), + "luma": PollSpec( + name="luma", + id_paths=("id",), + poll_url="/proxy/luma/generations/{id}", + status_paths=("state",), + success_values=("completed",), + failure_values=("failed",), + ), + "minimax": PollSpec( + name="minimax", + id_paths=("task_id",), + poll_url="/proxy/minimax/query/video_generation?task_id={id}", + status_paths=("status",), + success_values=("Success",), + failure_values=("Fail",), + post_success_url="/proxy/minimax/files/retrieve?file_id={id}", + post_success_id_paths=("file_id",), + ), + "runway": PollSpec( + name="runway", + id_paths=("id",), + poll_url="/proxy/runway/tasks/{id}", + status_paths=("status",), + success_values=("SUCCEEDED",), + failure_values=("FAILED", "CANCELLED", "THROTTLED"), + progress_path="progress", + ), + "moonvalley": PollSpec( + name="moonvalley", + id_paths=("id",), + poll_url="/proxy/moonvalley/prompts/{id}", + status_paths=("status",), + success_values=("completed",), + failure_values=("failed", "error"), + ), + "pika": PollSpec( + name="pika", + id_paths=("video_id", "id"), + poll_url="/proxy/pika/videos/{id}", + status_paths=("status",), + success_values=("finished",), + # Pika's enum has no explicit failure state; treat sustained queued/started + # as in-progress and rely on `timeout` to surface stalls. + failure_values=(), + progress_path="progress", + ), + "vidu": PollSpec( + name="vidu", + id_paths=("task_id",), + poll_url="/proxy/vidu/tasks/{id}/creations", + status_paths=("state",), + success_values=("success",), + failure_values=("failed",), + ), + "xai_video": PollSpec( + name="xai_video", + id_paths=("request_id",), + poll_url="/proxy/xai/v1/videos/{id}", + status_paths=("status",), + success_values=("done",), + failure_values=(), + ), } +def _build_poll_url(spec: PollSpec, job_id: str, create_path: str | None) -> str: + url = spec.poll_url.replace("{id}", str(job_id)) + if "{create_path}" in url: + if not create_path: + raise client.ApiError(0, "", f"{spec.name} poller needs the create path") + url = url.replace("{create_path}", create_path) + return url + + +def poll_generic( + initial: dict[str, Any], + api_key: str, + *, + spec: PollSpec, + create_path: str | None = None, + interval: float = 2.0, + timeout: float = 300.0, + on_progress: Callable[[float], None] | None = None, +) -> PollResult: + """Drive a partner's poll endpoint by reading dot-paths out of the JSON. + + ``initial`` is the create-response body; we pull a job id out of it, build + the poll URL from ``spec.poll_url``, and GET it until the status field hits + a terminal value. Handles MiniMax-style two-stage flows via + ``spec.post_success_url`` (a follow-up GET keyed off something the terminal + poll body contains, e.g. ``file_id``).""" + job_id = _first(initial, spec.id_paths) + if job_id is None: + raise client.ApiError(0, "", f"{spec.name} response missing id (looked for {spec.id_paths})") + url = _build_poll_url(spec, str(job_id), create_path) + deadline = _now() + timeout + last_body: dict[str, Any] = {} + while _now() < deadline: + resp = client.get(url, api_key=api_key) + if resp.status_code >= 400: + client.raise_for_status(resp) + last_body = resp.json() + if on_progress is not None and spec.progress_path: + p = _dotget(last_body, spec.progress_path) + if isinstance(p, int | float): + # Some partners report 0–100, others 0–1; normalize. + on_progress(float(p) / 100.0 if p > 1 else float(p)) + status = _first(last_body, spec.status_paths) + status_str = str(status) if status is not None else "" + if status_str in spec.success_values: + merged = dict(last_body) + if spec.post_success_url: + redeem_id = _first(last_body, spec.post_success_id_paths) + if redeem_id is not None: + redeem_url = spec.post_success_url.replace("{id}", str(redeem_id)) + r2 = client.get(redeem_url, api_key=api_key) + if r2.status_code < 400: + try: + merged["_redeemed"] = r2.json() + except ValueError: + pass + urls = _extract_urls(merged) + return PollResult(status="succeeded", raw=merged, image_urls=urls) + if status_str in spec.failure_values: + return PollResult(status="failed", raw=last_body, image_urls=[], error=status_str) + _sleep(interval) + return PollResult(status="failed", raw=last_body, image_urls=[], error=f"timed out after {timeout:.0f}s") + + +def extract_job_id(name: str, body: dict[str, Any]) -> str | None: + """Pull the partner's job id out of a create-response body for display.""" + if name == "bfl": + return body.get("id") or None + spec = _POLL_SPECS.get(name) + if spec is None: + return None + v = _first(body, spec.id_paths) + return str(v) if v is not None else None + + +def build_synthetic_initial(name: str, job_id: str, base_url: str | None = None) -> dict[str, Any]: + """Recreate a minimal create-response so ``poll_generic`` can find the id. + + Used by ``comfy generate resume`` β€” the user supplies just a partner key + and a job id, and we reverse-engineer the shape the poller expects.""" + if name == "bfl": + if not base_url: + raise client.ApiError(0, "", "BFL resume needs a base URL to build the polling_url") + return {"polling_url": f"{base_url}/proxy/bfl/get_result?id={job_id}"} + spec = _POLL_SPECS.get(name) + if not spec: + raise client.ApiError(0, "", f"No polling adapter for partner {name!r}") + primary = spec.id_paths[0] + body: dict[str, Any] = {} + cur = body + parts = primary.split(".") + for p in parts[:-1]: + cur[p] = {} + cur = cur[p] + cur[parts[-1]] = job_id + return body + + def get_poller(name: str) -> Callable[..., PollResult]: - try: - return _POLLERS[name] - except KeyError as e: - raise client.ApiError(0, "", f"No polling adapter for partner {name!r}") from e + """Return the poller callable for a partner name. + + All pollers accept the same kwargs (``api_key``, ``timeout``, ``on_progress``, + ``create_path``) so callers don't need to special-case which one they got.""" + if name == "bfl": + return poll_bfl + if name in _POLL_SPECS: + spec = _POLL_SPECS[name] + + def runner( + initial: dict[str, Any], + api_key: str, + *, + create_path: str | None = None, + interval: float = 2.0, + timeout: float = 300.0, + on_progress: Callable[[float], None] | None = None, + ) -> PollResult: + return poll_generic( + initial, + api_key, + spec=spec, + create_path=create_path, + interval=interval, + timeout=timeout, + on_progress=on_progress, + ) + + return runner + raise client.ApiError(0, "", f"No polling adapter for partner {name!r}") def sync_result_from_response(resp: httpx.Response) -> PollResult: """Wrap a sync response in a PollResult so the run path is uniform.""" - if resp.headers.get("content-type", "").startswith("image/"): + ctype = resp.headers.get("content-type", "") + if ctype.startswith(("image/", "video/", "audio/")): return PollResult(status="succeeded", raw={"_binary": True}, image_urls=[]) try: body = resp.json() diff --git a/comfy_cli/command/generate/spec.py b/comfy_cli/command/generate/spec.py index d9b6abfd..f9e0e3f1 100644 --- a/comfy_cli/command/generate/spec.py +++ b/comfy_cli/command/generate/spec.py @@ -11,6 +11,7 @@ from __future__ import annotations import os +import re as _re import time from dataclasses import dataclass from functools import lru_cache @@ -19,10 +20,27 @@ import yaml -try: - from yaml import CSafeLoader as _YamlLoader -except ImportError: - from yaml import SafeLoader as _YamlLoader # type: ignore[assignment] + +class _YamlLoader(yaml.SafeLoader): + """SafeLoader that strips YAML 1.1's bool aliases for ``on``/``off``/ + ``yes``/``no``/``y``/``n``. + + The vendored openapi uses unquoted ``[on, off]`` and similar as **string** + enum values (e.g. Kling's ``sound`` field), but PyYAML's default resolvers + promote them to ``True``/``False`` β€” which then breaks our flag-rendering + (`'|'.join` on a list with booleans) and the upstream API contract. Limit + bool resolution to the YAML 1.2 spelling (``true``/``false`` only).""" + + +_YamlLoader.yaml_implicit_resolvers = { + k: [(t, r) for (t, r) in resolvers if t != "tag:yaml.org,2002:bool"] + for k, resolvers in yaml.SafeLoader.yaml_implicit_resolvers.items() +} +_YamlLoader.add_implicit_resolver( + "tag:yaml.org,2002:bool", + _re.compile(r"^(?:true|True|TRUE|false|False|FALSE)$"), + list("tTfF"), +) PROXY_PREFIX = "/proxy/" DEFAULT_BASE_URL = "https://api.comfy.org" @@ -94,6 +112,30 @@ class Endpoint: "reve-edit": "reve/v1/image/edit", # Runway "runway": "runway/text_to_image", + # Video β€” Kling + "kling": "kling/v1/videos/text2video", + "kling-i2v": "kling/v1/videos/image2video", + "kling-extend": "kling/v1/videos/video-extend", + "kling-lipsync": "kling/v1/videos/lip-sync", + # Video β€” Luma Dream Machine + "luma": "luma/generations", + "luma-i2v": "luma/generations/image", + # Video β€” MiniMax / Hailuo + "hailuo": "minimax/video_generation", + # Video β€” Runway Gen-3 + "runway-i2v": "runway/image_to_video", + # Video β€” Moonvalley + "moonvalley-t2v": "moonvalley/prompts/text-to-video", + "moonvalley-i2v": "moonvalley/prompts/image-to-video", + # Video β€” Pika + "pika": "pika/generate/2.2/t2v", + "pika-i2v": "pika/generate/2.2/i2v", + # Video β€” Vidu + "vidu": "vidu/text2video", + "vidu-i2v": "vidu/img2video", + "vidu-extend": "vidu/extend", + # Video β€” xAI Grok + "grok-video": "xai/v1/videos/generations", } _PREFERRED_ALIAS: dict[str, str] = {v: k for k, v in _ALIASES.items()} @@ -119,9 +161,10 @@ def resolve_alias(target: str) -> str: return target -# Curated v1 image allowlist. Tuples of (endpoint_id, category, polling). +# Curated endpoint allowlist. Tuples of (endpoint_id, category, polling). +# ``polling`` is the partner-key the poll registry uses (None = sync). # Endpoint id is the openapi path with /proxy/ stripped. -_IMAGE_ALLOWLIST: list[tuple[str, str, str | None]] = [ +_ENDPOINT_ALLOWLIST: list[tuple[str, str, str | None]] = [ # OpenAI ("openai/images/generations", "text-to-image", None), ("openai/images/edits", "image-edit", None), @@ -162,8 +205,32 @@ def resolve_alias(target: str) -> str: # Reve ("reve/v1/image/create", "text-to-image", None), ("reve/v1/image/edit", "image-edit", None), - # Runway + # Runway (image) ("runway/text_to_image", "text-to-image", None), + # Video β€” Kling + ("kling/v1/videos/text2video", "text-to-video", "kling"), + ("kling/v1/videos/image2video", "image-to-video", "kling"), + ("kling/v1/videos/video-extend", "video-extend", "kling"), + ("kling/v1/videos/lip-sync", "lipsync", "kling"), + # Video β€” Luma + ("luma/generations", "text-to-video", "luma"), + ("luma/generations/image", "image-to-video", "luma"), + # Video β€” MiniMax / Hailuo + ("minimax/video_generation", "text-to-video", "minimax"), + # Video β€” Runway + ("runway/image_to_video", "image-to-video", "runway"), + # Video β€” Moonvalley + ("moonvalley/prompts/text-to-video", "text-to-video", "moonvalley"), + ("moonvalley/prompts/image-to-video", "image-to-video", "moonvalley"), + # Video β€” Pika + ("pika/generate/2.2/t2v", "text-to-video", "pika"), + ("pika/generate/2.2/i2v", "image-to-video", "pika"), + # Video β€” Vidu + ("vidu/text2video", "text-to-video", "vidu"), + ("vidu/img2video", "image-to-video", "vidu"), + ("vidu/extend", "video-extend", "vidu"), + # Video β€” xAI Grok + ("xai/v1/videos/generations", "text-to-video", "xai_video"), ] @@ -241,7 +308,7 @@ def _registry() -> dict[str, Endpoint]: spec = load_raw_spec() paths = spec.get("paths") or {} registry: dict[str, Endpoint] = {} - for endpoint_id, category, polling_hint in _IMAGE_ALLOWLIST: + for endpoint_id, category, polling_hint in _ENDPOINT_ALLOWLIST: path = PROXY_PREFIX + endpoint_id node = paths.get(path) if not node: diff --git a/tests/comfy_cli/command/generate/test_app.py b/tests/comfy_cli/command/generate/test_app.py index 2f35aea5..b82f618a 100644 --- a/tests/comfy_cli/command/generate/test_app.py +++ b/tests/comfy_cli/command/generate/test_app.py @@ -578,6 +578,120 @@ def test_generate_auto_upload_skipped_for_multipart(runner, api_key, tmp_path, m assert upload_called["hit"] is False +# ─── video models (async polling, generic poller path) ───────────────── + + +def test_video_kling_async_path(runner, api_key, monkeypatch): + """End-to-end async path through the generic kling poller.""" + submit = httpx.Response(200, json={"data": {"task_id": "k-xyz"}}) + finished = httpx.Response( + 200, + json={ + "data": { + "task_status": "succeed", + "task_result": {"videos": [{"url": "https://cdn.example/k.mp4"}]}, + } + }, + ) + monkeypatch.setattr(gen_app.client.httpx, "post", lambda *a, **kw: submit) + monkeypatch.setattr("comfy_cli.command.generate.client.get", lambda *a, **kw: finished) + monkeypatch.setattr("comfy_cli.command.generate.poll._sleep", lambda *_: None) + + r = runner.invoke(cli_app, ["generate", "kling", "--prompt", "a cat", "--duration", "5"]) + assert r.exit_code == 0, r.stdout + assert "https://cdn.example/k.mp4" in r.stdout + + +def test_video_luma_async_path(runner, api_key, monkeypatch): + submit = httpx.Response(200, json={"id": "luma-1", "state": "queued"}) + done = httpx.Response(200, json={"id": "luma-1", "state": "completed", "assets": {"video": "https://cdn/l.mp4"}}) + monkeypatch.setattr(gen_app.client.httpx, "post", lambda *a, **kw: submit) + monkeypatch.setattr("comfy_cli.command.generate.client.get", lambda *a, **kw: done) + monkeypatch.setattr("comfy_cli.command.generate.poll._sleep", lambda *_: None) + + r = runner.invoke( + cli_app, + [ + "generate", + "luma", + "--prompt", + "a cat", + "--aspect_ratio", + "16:9", + "--model", + "ray-2", + "--resolution", + "{}", + "--duration", + "{}", + ], + ) + assert r.exit_code == 0, r.stdout + assert "https://cdn/l.mp4" in r.stdout + + +def test_video_runway_failure_surfaces(runner, api_key, monkeypatch): + submit = httpx.Response(200, json={"id": "rw-1"}) + fail = httpx.Response(200, json={"id": "rw-1", "status": "FAILED"}) + monkeypatch.setattr(gen_app.client.httpx, "post", lambda *a, **kw: submit) + monkeypatch.setattr("comfy_cli.command.generate.client.get", lambda *a, **kw: fail) + monkeypatch.setattr("comfy_cli.command.generate.poll._sleep", lambda *_: None) + + r = runner.invoke( + cli_app, + [ + "generate", + "runway-i2v", + "--promptImage", + '"https://x/img.png"', + "--seed", + "1", + "--model", + "gen4_turbo", + "--duration", + "5", + "--ratio", + "1280:720", + ], + ) + assert r.exit_code == 1 + assert "FAILED" in r.stdout + + +def test_video_async_submission_shows_resume_alias(runner, api_key, monkeypatch): + submit = httpx.Response(200, json={"data": {"task_id": "k-async-1"}}) + monkeypatch.setattr(gen_app.client.httpx, "post", lambda *a, **kw: submit) + r = runner.invoke(cli_app, ["generate", "kling", "--prompt", "x", "--async"]) + assert r.exit_code == 0, r.stdout + assert "k-async-1" in r.stdout + assert "comfy generate resume kling k-async-1" in r.stdout + + +def test_video_resume_kling(runner, api_key, monkeypatch): + done = httpx.Response( + 200, + json={ + "data": { + "task_status": "succeed", + "task_result": {"videos": [{"url": "https://cdn/resumed.mp4"}]}, + } + }, + ) + monkeypatch.setattr("comfy_cli.command.generate.client.get", lambda *a, **kw: done) + monkeypatch.setattr("comfy_cli.command.generate.poll._sleep", lambda *_: None) + r = runner.invoke(cli_app, ["generate", "resume", "kling", "k-async-1"]) + assert r.exit_code == 0, r.stdout + assert "https://cdn/resumed.mp4" in r.stdout + + +def test_list_video_filter(runner): + r = runner.invoke(cli_app, ["generate", "list", "--style", "text-to-video"]) + assert r.exit_code == 0 + assert "kling" in r.stdout + assert "luma" in r.stdout + assert "pika" in r.stdout + + # ─── helpers: _arg_value / _separate_meta_flags ────────────────────────── diff --git a/tests/comfy_cli/command/generate/test_video_poll.py b/tests/comfy_cli/command/generate/test_video_poll.py new file mode 100644 index 00000000..5f8d4ac7 --- /dev/null +++ b/tests/comfy_cli/command/generate/test_video_poll.py @@ -0,0 +1,202 @@ +"""Tests for the generic config-driven poller and per-partner specs.""" + +import httpx +import pytest + +from comfy_cli.command.generate import poll + + +def _resp(body): + return httpx.Response(200, json=body) + + +@pytest.fixture +def no_sleep(monkeypatch): + monkeypatch.setattr(poll, "_sleep", lambda *_: None) + + +def _make_runner(get_responses): + """Patch client.get with an iterator over fake responses.""" + it = iter(get_responses) + return lambda *_a, **_kw: next(it) + + +def test_kling_sibling_poll_path(no_sleep, monkeypatch): + """Kling builds the poll URL from {create_path}/{id}.""" + captured = {} + + def fake_get(url, **kw): + captured["url"] = url + return _resp({"data": {"task_status": "succeed", "task_result": {"videos": [{"url": "https://cdn/v.mp4"}]}}}) + + monkeypatch.setattr("comfy_cli.command.generate.client.get", fake_get) + poller = poll.get_poller("kling") + result = poller( + {"data": {"task_id": "abc"}}, + api_key="comfyui-test", + create_path="/proxy/kling/v1/videos/text2video", + ) + assert captured["url"] == "/proxy/kling/v1/videos/text2video/abc" + assert result.status == "succeeded" + assert result.image_urls == ["https://cdn/v.mp4"] + + +def test_luma_succeeds(no_sleep, monkeypatch): + monkeypatch.setattr( + "comfy_cli.command.generate.client.get", + _make_runner( + [ + _resp({"id": "luma-1", "state": "dreaming"}), + _resp({"id": "luma-1", "state": "completed", "assets": {"video": "https://cdn/x.mp4"}}), + ] + ), + ) + result = poll.get_poller("luma")({"id": "luma-1", "state": "queued"}, api_key="k") + assert result.status == "succeeded" + assert "https://cdn/x.mp4" in result.image_urls + + +def test_runway_progress_normalized(no_sleep, monkeypatch): + """Runway reports progress as 0–1 floats β€” the poller forwards as-is.""" + seen: list[float] = [] + monkeypatch.setattr( + "comfy_cli.command.generate.client.get", + _make_runner( + [ + _resp({"id": "x", "status": "RUNNING", "progress": 0.5}), + _resp({"id": "x", "status": "SUCCEEDED", "output": ["https://cdn/v.mp4"]}), + ] + ), + ) + poll.get_poller("runway")({"id": "x"}, api_key="k", on_progress=seen.append) + assert seen == [0.5] + + +def test_runway_failure_states(no_sleep, monkeypatch): + monkeypatch.setattr( + "comfy_cli.command.generate.client.get", + _make_runner([_resp({"id": "x", "status": "CANCELLED"})]), + ) + result = poll.get_poller("runway")({"id": "x"}, api_key="k") + assert result.status == "failed" + assert "CANCELLED" in (result.error or "") + + +def test_minimax_redeems_file_id(no_sleep, monkeypatch): + """After Success, minimax needs a second GET to /files/retrieve to get the download URL.""" + monkeypatch.setattr( + "comfy_cli.command.generate.client.get", + _make_runner( + [ + _resp({"status": "Processing", "task_id": "t1"}), + _resp({"status": "Success", "task_id": "t1", "file_id": "f-42"}), + _resp({"file": {"download_url": "https://cdn/minimax.mp4"}}), + ] + ), + ) + result = poll.get_poller("minimax")({"task_id": "t1"}, api_key="k") + assert result.status == "succeeded" + assert "https://cdn/minimax.mp4" in result.image_urls + + +def test_pika_polls_videos_endpoint(no_sleep, monkeypatch): + captured = {} + + def fake_get(url, **kw): + captured["url"] = url + return _resp({"id": "v1", "status": "finished", "url": "https://cdn/p.mp4"}) + + monkeypatch.setattr("comfy_cli.command.generate.client.get", fake_get) + result = poll.get_poller("pika")({"video_id": "v1"}, api_key="k") + assert captured["url"] == "/proxy/pika/videos/v1" + assert result.status == "succeeded" + + +def test_vidu_polls_creations_path(no_sleep, monkeypatch): + captured = {} + + def fake_get(url, **kw): + captured["url"] = url + return _resp({"state": "success", "creations": [{"url": "https://cdn/vidu.mp4"}]}) + + monkeypatch.setattr("comfy_cli.command.generate.client.get", fake_get) + poll.get_poller("vidu")({"task_id": "t1"}, api_key="k") + assert captured["url"] == "/proxy/vidu/tasks/t1/creations" + + +def test_xai_video_polls_request_id(no_sleep, monkeypatch): + captured = {} + + def fake_get(url, **kw): + captured["url"] = url + return _resp({"status": "done", "video": {"url": "https://cdn/x.mp4"}}) + + monkeypatch.setattr("comfy_cli.command.generate.client.get", fake_get) + poll.get_poller("xai_video")({"request_id": "req-1"}, api_key="k") + assert captured["url"] == "/proxy/xai/v1/videos/req-1" + + +def test_moonvalley_polls_prompts(no_sleep, monkeypatch): + captured = {} + + def fake_get(url, **kw): + captured["url"] = url + return _resp({"id": "p-1", "status": "completed", "output_url": "https://cdn/m.mp4"}) + + monkeypatch.setattr("comfy_cli.command.generate.client.get", fake_get) + poll.get_poller("moonvalley")({"id": "p-1"}, api_key="k") + assert captured["url"] == "/proxy/moonvalley/prompts/p-1" + + +def test_missing_id_raises(monkeypatch): + monkeypatch.setattr("comfy_cli.command.generate.client.get", lambda *a, **kw: _resp({})) + with pytest.raises(Exception, match="missing id"): + poll.get_poller("kling")({}, api_key="k", create_path="/x") + + +def test_kling_without_create_path_raises(): + with pytest.raises(Exception, match="create path"): + poll.get_poller("kling")({"data": {"task_id": "abc"}}, api_key="k") + + +def test_build_synthetic_initial_for_each_partner(): + """Sanity-check the resume helper for every registered partner.""" + for name in ("kling", "luma", "minimax", "runway", "moonvalley", "pika", "vidu", "xai_video"): + body = poll.build_synthetic_initial(name, "abc") + assert poll.extract_job_id(name, body) == "abc", name + + +def test_build_synthetic_initial_for_bfl(): + body = poll.build_synthetic_initial("bfl", "abc", base_url="https://api.comfy.org") + assert "polling_url" in body + assert "abc" in body["polling_url"] + + +def test_extract_urls_recognizes_video_extensions(): + found = poll._extract_urls({"video_url": "https://cdn/x.mp4", "ignore": "https://cdn/notmedia"}) + assert "https://cdn/x.mp4" in found + assert "https://cdn/notmedia" not in found + + +def test_extract_urls_recognizes_query_strings(): + """Signed URLs with ?Expires=… shouldn't be excluded by their query string.""" + found = poll._extract_urls({"url": "https://cdn/v.mp4?Expires=123&Signature=abc"}) + assert found == ["https://cdn/v.mp4?Expires=123&Signature=abc"] + + +def test_extract_job_id_from_nested_paths(): + assert poll.extract_job_id("kling", {"data": {"task_id": "k1"}}) == "k1" + assert poll.extract_job_id("luma", {"id": "l1"}) == "l1" + assert poll.extract_job_id("minimax", {"task_id": "m1"}) == "m1" + assert poll.extract_job_id("xai_video", {"request_id": "x1"}) == "x1" + + +def test_existing_bfl_poller_still_works(no_sleep, monkeypatch): + """Regression: the original BFL adapter shouldn't be disturbed by the refactor.""" + monkeypatch.setattr( + "comfy_cli.command.generate.client.get", + _make_runner([_resp({"status": "Ready", "result": {"sample": "https://cdn/b.png"}})]), + ) + result = poll.get_poller("bfl")({"polling_url": "https://x"}, api_key="k") + assert result.status == "succeeded" + assert "https://cdn/b.png" in result.image_urls