diff --git a/.github/workflows/docker-streamdiffusion.yaml b/.github/workflows/docker-streamdiffusion.yaml index 09f70c0a..8729f80d 100644 --- a/.github/workflows/docker-streamdiffusion.yaml +++ b/.github/workflows/docker-streamdiffusion.yaml @@ -17,7 +17,6 @@ concurrency: jobs: base-streamdiffusion: name: comfyui-base:streamdiffusion image - if: ${{ github.repository == 'livepeer/comfystream' }} outputs: repository: ${{ steps.repo.outputs.repository }} image-digest: ${{ steps.build.outputs.digest }} @@ -37,7 +36,7 @@ jobs: id: repo shell: bash run: | - echo "repository=livepeer/comfyui-base" >> "$GITHUB_OUTPUT" + echo "repository=${{ github.repository_owner }}/comfyui-base" >> "$GITHUB_OUTPUT" echo "image-tag=streamdiffusion" >> "$GITHUB_OUTPUT" - name: Extract metadata (tags, labels) for Docker @@ -77,13 +76,12 @@ jobs: CACHEBUST=${{ github.run_id }} labels: ${{ steps.meta.outputs.labels }} annotations: ${{ steps.meta.outputs.annotations }} - cache-from: type=registry,ref=livepeer/comfyui-base:streamdiffusion-build-cache - cache-to: type=registry,mode=max,ref=livepeer/comfyui-base:streamdiffusion-build-cache + cache-from: type=registry,ref=${{ github.repository_owner }}/comfyui-base:streamdiffusion-build-cache + cache-to: type=registry,mode=max,ref=${{ github.repository_owner }}/comfyui-base:streamdiffusion-build-cache comfystream-streamdiffusion: name: comfystream:streamdiffusion image needs: base-streamdiffusion - if: ${{ github.repository == 'livepeer/comfystream' }} permissions: packages: write contents: read @@ -100,7 +98,7 @@ jobs: uses: docker/metadata-action@v5 with: images: | - livepeer/comfystream + ${{ github.repository }} tags: | type=raw,value=streamdiffusion type=raw,value=streamdiffusion-sha-${{ github.sha }} @@ -129,5 +127,5 @@ jobs: BASE_IMAGE=${{ needs.base-streamdiffusion.outputs.repository }}@${{ needs.base-streamdiffusion.outputs.image-digest }} labels: ${{ steps.meta.outputs.labels }} annotations: ${{ steps.meta.outputs.annotations }} - cache-from: type=registry,ref=livepeer/comfystream:streamdiffusion-build-cache - cache-to: type=registry,mode=max,ref=livepeer/comfystream:streamdiffusion-build-cache + cache-from: type=registry,ref=${{ github.repository }}:streamdiffusion-build-cache + cache-to: type=registry,mode=max,ref=${{ github.repository }}:streamdiffusion-build-cache diff --git a/.github/workflows/docker.yaml b/.github/workflows/docker.yaml index 1d79dfe0..cbc90790 100644 --- a/.github/workflows/docker.yaml +++ b/.github/workflows/docker.yaml @@ -17,7 +17,6 @@ concurrency: jobs: base: name: comfyui-base image - if: ${{ github.repository == 'livepeer/comfystream' }} outputs: repository: ${{ steps.repo.outputs.repository }} image-digest: ${{ steps.build.outputs.digest }} @@ -36,7 +35,7 @@ jobs: id: repo shell: bash run: | - echo "repository=livepeer/comfyui-base" >> "$GITHUB_OUTPUT" + echo "repository=${{ github.repository_owner }}/comfyui-base" >> "$GITHUB_OUTPUT" - name: Extract metadata (tags, labels) for Docker id: meta @@ -80,13 +79,12 @@ jobs: CACHEBUST=${{ github.run_id }} labels: ${{ steps.meta.outputs.labels }} annotations: ${{ steps.meta.outputs.annotations }} - cache-from: type=registry,ref=livepeer/comfyui-base:build-cache - cache-to: type=registry,mode=max,ref=livepeer/comfyui-base:build-cache + cache-from: type=registry,ref=${{ github.repository_owner }}/comfyui-base:build-cache + cache-to: type=registry,mode=max,ref=${{ github.repository_owner }}/comfyui-base:build-cache comfystream: name: comfystream image needs: base - if: ${{ github.repository == 'livepeer/comfystream' }} permissions: packages: write contents: read diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 98baa6cf..91c251d7 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -15,7 +15,6 @@ concurrency: jobs: codeql: name: Perform CodeQL analysis - if: ${{ github.repository == 'livepeer/comfystream' }} runs-on: ubuntu-latest steps: - name: Check out code @@ -68,7 +67,7 @@ jobs: - name: Set up Python uses: actions/setup-python@v6 with: - python-version: '3.12' + python-version: '3.12.x' cache: pip - name: Install dependencies diff --git a/.vscode/launch.json b/.vscode/launch.json index f05e02f5..45707a9a 100755 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -41,7 +41,7 @@ "name": "Run ComfyStream BYOC", "type": "debugpy", "request": "launch", - "cwd": "/workspace/ComfyUI", + "cwd": "/workspace/comfystream", "program": "/workspace/comfystream/server/byoc.py", "console": "integratedTerminal", "args": [ diff --git a/README.md b/README.md index a9edf4be..45ff597c 100644 --- a/README.md +++ b/README.md @@ -20,6 +20,7 @@ This repo also includes a WebRTC server and UI that uses comfystream to support - [Run UI](#run-ui) - [Limitations](#limitations) - [Troubleshoot](#troubleshoot) + - [Livepeer live-runner](#livepeer-live-runner) ## Quick Start @@ -160,6 +161,8 @@ If you only have a subset of those UDP ports available, you can use the `--media python server/app.py --workspace --media-ports 1024,1025,... ``` +> Tip: Use `--workspace` (preferred). `--cwd` remains a compatible alias and honors `COMFYUI_CWD`. + If you are running the server in a restrictive network environment where this is not possible, you will need to use a TURN server. At the moment, the server supports using Twilio's TURN servers (although it is easy to make the update to support arbitrary TURN servers): @@ -224,3 +227,35 @@ This project has been tested locally successfully with the following setup: - Driver: 550.127.05 - CUDA: 12.5 - torch: 2.5.1+cu121 + +## Livepeer live-runner + +Register ComfyStream against a go-livepeer orchestrator with `-useLiveRunners` and drive it through the SDK (same path as the transcode live-runner). One process, capacity **1**, metered by session wall-clock — sell latency/liveness, not a batch $/image race. + +Agent-shaped endpoints: + +| Method | Path | Role | +| --- | --- | --- | +| `POST` | `/analyze` | Video-in → text-out (build/demo this first) | +| `POST` | `/start_stream` | Live trickle video (optional text channel) | +| `POST` | `/update_stream` | Mid-session prompt / resolution update | +| `GET` | `/text` | Buffered text outputs for the active session | +| `GET` | `/healthz` | Health | + +Attach to an already-running orchestrator (example targets ai1): + +```sh +docker compose -f docker-compose.live-runner.yml up -d --build +curl -sk https://ai1.eliteencoder.net:8936/discovery | jq '.[].runners[].app' +``` + +Smoke client (after the runner appears in discovery): + +```sh +pip install "livepeer-gateway @ git+https://github.com/livepeer/livepeer-python-gateway@ja/live-runner" av aiohttp +python server/live_runner_client.py sample.mp4 \ + --workflow path/to/video-in-text-out.json \ + --discovery https://ai1.eliteencoder.net:8936/discovery +``` + +Optional dep: `pip install '.[live-runner]'`. Legacy BYOC (`server/byoc.py`) is unchanged. diff --git a/docker-compose.live-runner.yml b/docker-compose.live-runner.yml new file mode 100644 index 00000000..8657a38c --- /dev/null +++ b/docker-compose.live-runner.yml @@ -0,0 +1,43 @@ +# Attach ComfyStream as a live-runner to an already-running orchestrator. +# Does not start go-livepeer — use the same orch as live-runner-transcode. +# +# docker compose -f docker-compose.live-runner.yml up -d --build +# +# Defaults target ai1 (override via .env). Capacity stays 1; price is fair +# latency pricing, not a batch $/image race. + +services: + comfystream-live-runner: + build: + context: . + dockerfile: docker/Dockerfile.live-runner + image: comfystream-live-runner:latest + container_name: live-runner-comfystream + restart: unless-stopped + network_mode: host + environment: + - NVIDIA_VISIBLE_DEVICES=${COMFYSTREAM_GPU_UUID:-GPU-745244ef-3615-57bd-458f-299fe2289f5a} + - COMFYUI_CWD=/workspace/ComfyUI + - PYTHONUNBUFFERED=1 + volumes: + # Shared model/storage volume used by other AI workloads on this host. + - ${COMFYSTREAM_MODELS_DIR:-/livepeer/ai/data/models}:/workspace/ComfyUI/models + - ${COMFYSTREAM_STORAGE_DIR:-/livepeer/ai/data}:/app/storage + command: + - --live-runner + - --host=0.0.0.0 + - --port=${COMFYSTREAM_PORT:-8991} + - --orchestrator=${LIVEPEER_ORCH_URL:-https://ai1.eliteencoder.net:8936} + - --orchSecret=${LIVEPEER_ORCH_SECRET:-LlFOUpAP8uTDIuz} + - --runner-url=${LIVEPEER_RUNNER_URL:-https://ai1.eliteencoder.net:8991} + - --workspace=/workspace/ComfyUI + - --capacity=1 + - --price=${COMFYSTREAM_PRICE:-0.001} + deploy: + resources: + reservations: + devices: + - capabilities: [gpu] + driver: nvidia + device_ids: + - "${COMFYSTREAM_GPU_UUID:-GPU-745244ef-3615-57bd-458f-299fe2289f5a}" diff --git a/docker/Dockerfile.live-runner b/docker/Dockerfile.live-runner new file mode 100644 index 00000000..e636aa10 --- /dev/null +++ b/docker/Dockerfile.live-runner @@ -0,0 +1,24 @@ +# Overlay image: ComfyStream + livepeer-gateway live-runner entrypoint. +ARG BASE_IMAGE=livepeer/comfystream:latest +FROM ${BASE_IMAGE} + +WORKDIR /workspace/comfystream + +# Install the live-runner SDK into the ComfyStream conda env. +# livepeer-gateway's generated pb2 requires protobuf runtime >= gencode (6.31.1). +# ComfyUI pins protobuf<5, but the live-runner path still runs with protobuf 6. +RUN /bin/bash -lc 'source /workspace/miniconda3/etc/profile.d/conda.sh \ + && conda activate comfystream \ + && pip install --no-cache-dir \ + "protobuf>=6.31.1" \ + "livepeer-gateway @ git+https://github.com/livepeer/livepeer-python-gateway.git@ja/live-runner" \ + av' + +COPY server/live_runner.py /workspace/comfystream/server/live_runner.py +COPY docker/entrypoint.sh /workspace/comfystream/docker/entrypoint.sh +RUN chmod +x /workspace/comfystream/docker/entrypoint.sh + +EXPOSE 8991 + +ENTRYPOINT ["/workspace/comfystream/docker/entrypoint.sh"] +CMD ["--live-runner"] diff --git a/docker/entrypoint.sh b/docker/entrypoint.sh index f3101859..c85b12ed 100755 --- a/docker/entrypoint.sh +++ b/docker/entrypoint.sh @@ -15,6 +15,7 @@ show_help() { echo " --server Start ComfyUI only" echo " --api Start ComfyStream API Server only" echo " --ui Start ComfyStream UI only" + echo " --live-runner Start ComfyStream as a Livepeer live-runner" echo " --help Show this help message" echo "" } @@ -29,11 +30,13 @@ WORKSPACE_STORAGE="/app/storage" COMFYUI_DIR="/workspace/ComfyUI" MODELS_DIR="$COMFYUI_DIR/models" OUTPUT_DIR="$COMFYUI_DIR/output" +export COMFYUI_CWD="$COMFYUI_DIR" # Initialize variables to track which services to start START_COMFYUI=false START_API=false START_UI=false +START_LIVE_RUNNER=false # First pass: check for service flags and set variables for arg in "$@"; do @@ -47,6 +50,9 @@ for arg in "$@"; do --ui) START_UI=true ;; + --live-runner) + START_LIVE_RUNNER=true + ;; esac done @@ -231,6 +237,21 @@ fi cd /workspace/comfystream +# Live-runner path: register with go-livepeer -useLiveRunners and drive Pipeline +# in-process (analyze / start_stream / update_stream). Remaining args after the +# flag are forwarded to server/live_runner.py. +if [ "$START_LIVE_RUNNER" = true ]; then + conda activate comfystream + # Drop the --live-runner flag; pass the rest through. + shift_args=() + for arg in "$@"; do + if [ "$arg" != "--live-runner" ]; then + shift_args+=("$arg") + fi + done + exec python server/live_runner.py "${shift_args[@]}" +fi + # If any service flags were specified, start supervisord and the requested services if [ "$START_COMFYUI" = true ] || [ "$START_API" = true ] || [ "$START_UI" = true ]; then # Start supervisord in background diff --git a/install.py b/install.py index 7552ea8a..4fe3a04a 100644 --- a/install.py +++ b/install.py @@ -75,7 +75,9 @@ def download_and_extract_ui_files(version: str): parser = argparse.ArgumentParser(description="Install custom node requirements") parser.add_argument( "--workspace", - default=os.environ.get("COMFY_UI_WORKSPACE", None), + "--cwd", + dest="workspace", + default=os.environ.get("COMFYUI_CWD"), required=False, help="Set Comfy workspace", ) @@ -100,6 +102,9 @@ def download_and_extract_ui_files(version: str): break current = os.path.dirname(current) + if workspace is not None: + os.environ.setdefault("COMFYUI_CWD", workspace) + logger.info("Installing comfystream package...") subprocess.check_call([sys.executable, "-m", "pip", "install", "-e", "."]) diff --git a/pyproject.toml b/pyproject.toml index cb708f86..d8a5fc4a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,6 +24,10 @@ dev = ["pytest", "pytest-cov", "ruff"] server = [ "pytrickle @ git+https://github.com/livepeer/pytrickle.git@v0.1.8" ] +live-runner = [ + "livepeer-gateway @ git+https://github.com/livepeer/livepeer-python-gateway.git@ja/live-runner", + "av", +] [project.urls] Repository = "https://github.com/yondonfu/comfystream" diff --git a/server/app.py b/server/app.py index b93e35ee..7723a4c1 100644 --- a/server/app.py +++ b/server/app.py @@ -654,15 +654,29 @@ async def on_startup(app: web.Application): if app["media_ports"]: patch_loop_datagram(app["media_ports"]) + # Always set the workspace as cwd so ComfyUI resolves paths correctly. + comfy_kwargs = { + "cwd": app["workspace"], + } + if app.get("config"): + # Pass config directly to ComfyUI; avoid adding other ComfyUI flags, + # but still preserve cwd so config paths resolve relative to workspace. + comfy_kwargs["config"] = app["config"] + else: + comfy_kwargs.update( + { + "disable_cuda_malloc": True, + "gpu_only": True, + "preview_method": "none", + "logging_level": app.get("logging_level", None), + "blacklist_custom_nodes": ["ComfyUI-Manager"], + } + ) + app["pipeline"] = Pipeline( width=512, height=512, - cwd=app["workspace"], - disable_cuda_malloc=True, - gpu_only=True, - preview_method="none", - comfyui_inference_log_level=app.get("comfyui_inference_log_level", None), - blacklist_custom_nodes=["ComfyUI-Manager"], + **comfy_kwargs, ) await app["pipeline"].initialize() app["pcs"] = set() @@ -681,7 +695,21 @@ async def on_shutdown(app: web.Application): parser.add_argument("--port", default=8889, help="Set the signaling port") parser.add_argument("--media-ports", default=None, help="Set the UDP ports for WebRTC media") parser.add_argument("--host", default="127.0.0.1", help="Set the host") - parser.add_argument("--workspace", default=None, required=True, help="Set Comfy workspace") + parser.add_argument( + "-c", + "--config", + dest="config", + default=None, + help="Path to ComfyUI config file (yaml/json/ini/conf). When provided, it is passed directly to ComfyUI.", + ) + parser.add_argument( + "--workspace", + "--cwd", + dest="workspace", + default=os.environ.get("COMFYUI_CWD"), + required=True, + help="Set ComfyUI workspace directory (preferred; alias: --cwd)", + ) parser.add_argument( "--log-level", default="INFO", @@ -707,10 +735,10 @@ async def on_shutdown(app: web.Application): help="Set the global logging level for ComfyUI", ) parser.add_argument( - "--comfyui-inference-log-level", + "--logging-level", default=None, choices=logging._nameToLevel.keys(), - help="Set the logging level for ComfyUI inference", + help="Set the logging level for ComfyUI (passed through to ComfyUI Configuration)", ) args = parser.parse_args() @@ -720,9 +748,14 @@ async def on_shutdown(app: web.Application): datefmt="%H:%M:%S", ) + logger.info(f"Using ComfyUI workspace: {args.workspace}") + os.environ.setdefault("COMFYUI_CWD", args.workspace) + app = web.Application() app["media_ports"] = args.media_ports.split(",") if args.media_ports else None app["workspace"] = args.workspace + app["logging_level"] = args.logging_level + app["config"] = args.config app.on_startup.append(on_startup) app.on_shutdown.append(on_shutdown) @@ -768,7 +801,5 @@ def force_print(*args, **kwargs): timeout_filter = ComfyStreamTimeoutFilter() logging.getLogger("comfy.cmd.execution").addFilter(timeout_filter) logging.getLogger("comfystream").addFilter(timeout_filter) - if args.comfyui_inference_log_level: - app["comfyui_inference_log_level"] = args.comfyui_inference_log_level web.run_app(app, host=args.host, port=int(args.port), print=force_print) diff --git a/server/byoc.py b/server/byoc.py index 3f8f3470..43c82243 100644 --- a/server/byoc.py +++ b/server/byoc.py @@ -30,10 +30,19 @@ def main(): ) parser.add_argument("--port", default=8000, help="Set the server port") parser.add_argument("--host", default="0.0.0.0", help="Set the host") + parser.add_argument( + "-c", + "--config", + dest="config", + default=None, + help="Path to ComfyUI config file (yaml/json/ini/conf). When provided, it is passed directly to ComfyUI.", + ) parser.add_argument( "--workspace", - default=os.getcwd() + "/../ComfyUI", - help="Set Comfy workspace (Default: ../ComfyUI)", + "--cwd", + dest="workspace", + default=os.environ.get("COMFYUI_CWD", os.getcwd() + "/../ComfyUI"), + help="Set ComfyUI workspace directory (preferred; alias: --cwd)", ) parser.add_argument( "--log-level", @@ -48,10 +57,10 @@ def main(): help="Set the global logging level for ComfyUI", ) parser.add_argument( - "--comfyui-inference-log-level", + "--logging-level", default=None, choices=logging._nameToLevel.keys(), - help="Set the logging level for ComfyUI inference", + help="Set the logging level for ComfyUI (passed through to ComfyUI Configuration)", ) parser.add_argument( "--disable-frame-skip", @@ -80,6 +89,9 @@ def main(): ) logging.getLogger("comfy.model_detection").setLevel(logging.WARNING) + logger.info(f"Using ComfyUI workspace: {args.workspace}") + os.environ.setdefault("COMFYUI_CWD", args.workspace) + # Allow overriding of ComfyUI log levels. if args.comfyui_log_level: log_level = logging._nameToLevel.get(args.comfyui_log_level.upper()) @@ -104,13 +116,13 @@ def force_print(*args, **kwargs): frame_processor = ComfyStreamFrameProcessor( width=args.width, height=args.height, + config=args.config, workspace=args.workspace, disable_cuda_malloc=True, gpu_only=True, preview_method="none", blacklist_custom_nodes=["ComfyUI-Manager"], - logging_level=args.comfyui_log_level, - comfyui_inference_log_level=args.comfyui_inference_log_level, + logging_level=args.logging_level or args.comfyui_log_level, ) # Create frame skip configuration only if enabled diff --git a/server/frame_processor.py b/server/frame_processor.py index 3dcb835d..9de7651c 100644 --- a/server/frame_processor.py +++ b/server/frame_processor.py @@ -329,16 +329,29 @@ async def load_model(self, **kwargs): params = {**self._load_params, **kwargs} if self.pipeline is None: + # Always set the workspace as cwd so ComfyUI resolves paths correctly. + comfy_kwargs: Dict[str, Any] = { + "cwd": params.get("workspace", os.getcwd()), + } + + # If a ComfyUI config file is provided, prefer it and pass it through. + if params.get("config"): + # Pass config directly to ComfyUI; avoid adding other ComfyUI flags, + # but still preserve cwd so config paths resolve relative to workspace. + comfy_kwargs["config"] = params["config"] + else: + comfy_kwargs.update({ + "disable_cuda_malloc": params.get("disable_cuda_malloc", True), + "gpu_only": params.get("gpu_only", True), + "preview_method": params.get("preview_method", "none"), + "logging_level": params.get("logging_level", "INFO"), + "blacklist_custom_nodes": ["ComfyUI-Manager"], + }) + self.pipeline = Pipeline( width=int(params.get("width", 512)), height=int(params.get("height", 512)), - cwd=params.get("workspace", os.getcwd()), - disable_cuda_malloc=params.get("disable_cuda_malloc", True), - gpu_only=params.get("gpu_only", True), - preview_method=params.get("preview_method", "none"), - comfyui_inference_log_level=params.get("comfyui_inference_log_level", "INFO"), - logging_level=params.get("comfyui_inference_log_level", "INFO"), - blacklist_custom_nodes=["ComfyUI-Manager"], + **comfy_kwargs, ) await self.pipeline.initialize() diff --git a/server/live_runner.py b/server/live_runner.py new file mode 100644 index 00000000..47cf3913 --- /dev/null +++ b/server/live_runner.py @@ -0,0 +1,537 @@ +#!/usr/bin/env python3 +"""ComfyStream live-runner: Pipeline-backed analyze + live stream on Livepeer. + +Registers as app ``comfystream`` (persistent, capacity 1) and drives +``comfystream.Pipeline`` in-process — no BYOC/pytrickle subprocess. + +Agent surface: + POST /analyze video-in → text-out (build this first) + POST /start_stream live video (and optional text) trickle session + POST /update_stream mid-session prompt / resolution update + GET /text buffered text outputs for the active session + GET /healthz + +Livepeer integration (grep ``# Livepeer:``): + 1. register_runner() + 2. create_trickle_channels() + 3. registration.close() / on_session_release cleanup +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import logging +import os +from contextlib import suppress +from dataclasses import dataclass, field +from typing import Any, Optional + +from aiohttp import web + +from comfystream.modalities import WorkflowModality +from comfystream.pipeline import Pipeline +from comfystream.utils import convert_prompt +from livepeer_gateway.channel_writer import JSONLWriter +from livepeer_gateway.live_runner import register_runner +from livepeer_gateway.media_decode import AudioDecodedMediaFrame, VideoDecodedMediaFrame +from livepeer_gateway.media_output import MediaOutput +from livepeer_gateway.media_publish import MediaPublish + +log = logging.getLogger("comfystream-live-runner") + +APP_ID = "comfystream" +DEFAULT_HOST = "0.0.0.0" +DEFAULT_PORT = 8991 +CHANNEL_MIME_VIDEO = "video/mp2t" +CHANNEL_MIME_JSONL = "application/jsonl" +TEXT_POLL_INTERVAL = 0.25 + + +@dataclass +class RunnerSession: + session_id: str + kind: str # "analyze" | "stream" + io: WorkflowModality + in_url: str + out_url: str | None = None + text_url: str | None = None + media_in: MediaOutput | None = None + video_out: MediaPublish | None = None + text_out: JSONLWriter | None = None + text_task: asyncio.Task | None = None + collected_text: list[str] = field(default_factory=list) + prompts: Any = None + + def to_json(self) -> dict[str, Any]: + data: dict[str, Any] = { + "session": self.session_id, + "kind": self.kind, + "in": self.in_url, + "modalities": self.io, + } + if self.out_url: + data["out"] = self.out_url + if self.text_url: + data["text"] = self.text_url + return data + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="ComfyStream Livepeer live-runner.") + parser.add_argument("--orchestrator", default=os.environ.get("LIVEPEER_ORCH_URL", "https://localhost:8935")) + parser.add_argument( + "--orchSecret", + default=os.environ.get("LIVEPEER_ORCH_SECRET", os.environ.get("ORCH_SECRET", "abcdef")), + ) + parser.add_argument( + "--runner-url", + default=os.environ.get("LIVEPEER_RUNNER_URL", f"http://127.0.0.1:{DEFAULT_PORT}"), + ) + parser.add_argument("--host", default=os.environ.get("LIVEPEER_RUNNER_HOST", DEFAULT_HOST)) + parser.add_argument("--port", type=int, default=int(os.environ.get("LIVEPEER_RUNNER_PORT", str(DEFAULT_PORT)))) + parser.add_argument( + "--workspace", + default=os.environ.get("COMFYUI_CWD", os.environ.get("COMFYUI_WORKSPACE", "")), + help="ComfyUI workspace directory (COMFYUI_CWD).", + ) + parser.add_argument("--width", type=int, default=512) + parser.add_argument("--height", type=int, default=512) + parser.add_argument( + "--price", + type=float, + default=float(os.environ.get("LIVEPEER_RUNNER_PRICE", "0.001")), + help="USD per hour (metered). Fair latency pricing — not a batch $/image race.", + ) + parser.add_argument( + "--capacity", + type=int, + default=int(os.environ.get("LIVEPEER_RUNNER_CAPACITY", "1")), + ) + parser.add_argument( + "--skip-bootstrap", + action="store_true", + help="Skip Pipeline default-workflow bootstrap (faster bring-up if workspace is ready).", + ) + return parser.parse_args() + + +def _session_id(request: web.Request) -> str: + session_id = request.headers.get("Livepeer-Session-Id", "").strip() + if not session_id: + raise web.HTTPBadRequest(text="missing Livepeer-Session-Id header") + return session_id + + +def _channel_url(channel: dict[str, Any], *, internal: bool = False) -> str: + if internal: + return str(channel.get("internal_url") or channel["url"]) + return str(channel["url"]) + + +def _extract_prompts(payload: dict[str, Any]) -> Any: + prompts = payload.get("prompts", payload.get("prompt")) + if prompts is None: + raise web.HTTPBadRequest(text="missing prompts/prompt in body") + if isinstance(prompts, str): + prompts = json.loads(prompts) + return prompts + + +def _convert_prompts(prompts: Any) -> list[dict[str, Any]]: + if isinstance(prompts, list): + return [convert_prompt(p, return_dict=True) for p in prompts] + return [convert_prompt(prompts, return_dict=True)] + + +def _require_analyze_io(io: WorkflowModality) -> None: + if not io["video"]["input"]: + raise web.HTTPBadRequest(text="analyze requires a workflow with video input") + if not io["text"]["output"]: + raise web.HTTPBadRequest(text="analyze requires a workflow with text output") + + +def _require_stream_io(io: WorkflowModality) -> None: + if not (io["video"]["input"] or io["audio"]["input"] or io["video"]["output"]): + raise web.HTTPBadRequest(text="start_stream requires a workflow with media I/O") + + +async def _close_session(app: web.Application, *, stop_prompts: bool = True) -> None: + session: RunnerSession | None = app.get("session") + if session is None: + return + app["session"] = None + + if session.text_task is not None and not session.text_task.done(): + session.text_task.cancel() + with suppress(asyncio.CancelledError, Exception): + await session.text_task + + with suppress(Exception): + if session.media_in is not None: + await session.media_in.close() + with suppress(Exception): + if session.video_out is not None: + await session.video_out.close() + with suppress(Exception): + if session.text_out is not None: + await session.text_out.close() + + pipeline: Pipeline | None = app.get("pipeline") + if stop_prompts and pipeline is not None: + with suppress(Exception): + await pipeline.stop_prompts(cleanup=True) + + +async def _on_session_release(app: web.Application, event: Any) -> None: + session_id = getattr(event, "session_id", "") or "" + session: RunnerSession | None = app.get("session") + if session is None: + return + if session_id and session.session_id != session_id: + return + log.info("orchestrator released session %s; cleaning up", session.session_id) + await _close_session(app) + + +async def _text_forward_loop(app: web.Application, session: RunnerSession) -> None: + pipeline: Pipeline = app["pipeline"] + while True: + try: + text = await pipeline.get_text_output() + if text is None or str(text).strip() == "": + await asyncio.sleep(TEXT_POLL_INTERVAL) + continue + text_str = str(text) + session.collected_text.append(text_str) + if session.text_out is not None: + await session.text_out.write({"type": "text", "text": text_str}) + except asyncio.CancelledError: + raise + except Exception: + log.exception("text forwarder error") + await asyncio.sleep(TEXT_POLL_INTERVAL) + + +async def _handle_video_frame( + app: web.Application, + session: RunnerSession, + decoded: AudioDecodedMediaFrame | VideoDecodedMediaFrame, +) -> None: + if decoded.kind != "video": + return + pipeline: Pipeline = app["pipeline"] + frame = decoded.frame + await pipeline.put_video_frame(frame) + if pipeline.produces_video_output(): + out = await pipeline.get_processed_video_frame() + if session.video_out is not None: + await session.video_out.write_frame(out) + else: + # Video-in / text-out: drain the sync queue without waiting for a video tensor. + await pipeline.video_incoming_frames.get() + + +async def _apply_workflow( + pipeline: Pipeline, + prompts: Any, + *, + width: Optional[int], + height: Optional[int], + skip_warmup: bool = False, +) -> WorkflowModality: + converted = _convert_prompts(prompts) + if width and width > 0: + pipeline.width = int(width) + if height and height > 0: + pipeline.height = int(height) + await pipeline.apply_prompts(converted, skip_warmup=skip_warmup) + if not skip_warmup: + await pipeline.ensure_warmup(pipeline.width, pipeline.height) + if pipeline.state_manager.can_stream(): + await pipeline.start_streaming() + return pipeline.get_workflow_io_capabilities() + + +async def _handle_analyze(request: web.Request) -> web.Response: + app = request.app + session_id = _session_id(request) + existing: RunnerSession | None = app.get("session") + if existing is not None: + if existing.session_id != session_id: + raise web.HTTPConflict(text="runner already has an active session") + return web.json_response(existing.to_json()) + + payload = json.loads(await request.read() or b"{}") + if not isinstance(payload, dict): + raise web.HTTPBadRequest(text="body must be a JSON object") + prompts = _extract_prompts(payload) + width = payload.get("width") + height = payload.get("height") + + pipeline: Pipeline = app["pipeline"] + try: + io = await _apply_workflow( + pipeline, + prompts, + width=int(width) if width else None, + height=int(height) if height else None, + ) + except web.HTTPException: + raise + except Exception as exc: + log.exception("failed to apply analyze workflow") + raise web.HTTPBadRequest(text=f"invalid workflow: {exc}") from exc + + _require_analyze_io(io) + + channels = await app["registration"].create_trickle_channels( # Livepeer: 2 + request, + [ + {"name": "in", "mime_type": CHANNEL_MIME_VIDEO}, + {"name": "text", "mime_type": CHANNEL_MIME_JSONL}, + ], + ) + by_name = {c["name"]: c for c in channels} + if "in" not in by_name or "text" not in by_name: + raise web.HTTPInternalServerError(text="orchestrator did not return in/text channels") + + session = RunnerSession( + session_id=session_id, + kind="analyze", + io=io, + in_url=_channel_url(by_name["in"]), + text_url=_channel_url(by_name["text"]), + text_out=JSONLWriter(_channel_url(by_name["text"], internal=True)), + prompts=prompts, + ) + + async def _on_frame(decoded) -> None: + await _handle_video_frame(app, session, decoded) + + session.media_in = MediaOutput( + _channel_url(by_name["in"], internal=True), + on_frame=_on_frame, + ) + session.text_task = asyncio.create_task(_text_forward_loop(app, session)) + app["session"] = session + + for task in session.media_in.callback_tasks(): + task.add_done_callback( + lambda _t: asyncio.create_task(_close_session(app)) + ) + + log.info("started analyze session %s", session_id) + return web.json_response(session.to_json()) + + +async def _handle_start_stream(request: web.Request) -> web.Response: + app = request.app + session_id = _session_id(request) + existing: RunnerSession | None = app.get("session") + if existing is not None: + if existing.session_id != session_id: + raise web.HTTPConflict(text="runner already has an active session") + return web.json_response(existing.to_json()) + + payload = json.loads(await request.read() or b"{}") + if not isinstance(payload, dict): + raise web.HTTPBadRequest(text="body must be a JSON object") + prompts = _extract_prompts(payload) + width = payload.get("width") + height = payload.get("height") + + pipeline: Pipeline = app["pipeline"] + try: + io = await _apply_workflow( + pipeline, + prompts, + width=int(width) if width else None, + height=int(height) if height else None, + ) + except Exception as exc: + log.exception("failed to apply stream workflow") + raise web.HTTPBadRequest(text=f"invalid workflow: {exc}") from exc + + _require_stream_io(io) + + channel_reqs: list[dict[str, str]] = [] + if io["video"]["input"] or io["audio"]["input"]: + channel_reqs.append({"name": "in", "mime_type": CHANNEL_MIME_VIDEO}) + if io["video"]["output"]: + channel_reqs.append({"name": "out", "mime_type": CHANNEL_MIME_VIDEO}) + if io["text"]["output"]: + channel_reqs.append({"name": "text", "mime_type": CHANNEL_MIME_JSONL}) + if not channel_reqs: + raise web.HTTPBadRequest(text="workflow produced no trickle channels") + + channels = await app["registration"].create_trickle_channels( # Livepeer: 2 + request, + channel_reqs, + ) + by_name = {c["name"]: c for c in channels} + + session = RunnerSession( + session_id=session_id, + kind="stream", + io=io, + in_url=_channel_url(by_name["in"]) if "in" in by_name else "", + out_url=_channel_url(by_name["out"]) if "out" in by_name else None, + text_url=_channel_url(by_name["text"]) if "text" in by_name else None, + prompts=prompts, + ) + if "out" in by_name: + session.video_out = MediaPublish(_channel_url(by_name["out"], internal=True)) + if "text" in by_name: + session.text_out = JSONLWriter(_channel_url(by_name["text"], internal=True)) + session.text_task = asyncio.create_task(_text_forward_loop(app, session)) + + if "in" in by_name: + async def _on_frame(decoded) -> None: + await _handle_video_frame(app, session, decoded) + + session.media_in = MediaOutput( + _channel_url(by_name["in"], internal=True), + on_frame=_on_frame, + ) + for task in session.media_in.callback_tasks(): + task.add_done_callback( + lambda _t: asyncio.create_task(_close_session(app)) + ) + + app["session"] = session + log.info("started stream session %s", session_id) + return web.json_response(session.to_json()) + + +async def _handle_update_stream(request: web.Request) -> web.Response: + app = request.app + session_id = _session_id(request) + session: RunnerSession | None = app.get("session") + if session is None: + raise web.HTTPNotFound(text="no active session") + if session.session_id != session_id: + raise web.HTTPConflict(text="runner has a different active session") + + payload = json.loads(await request.read() or b"{}") + if not isinstance(payload, dict): + raise web.HTTPBadRequest(text="body must be a JSON object") + + pipeline: Pipeline = app["pipeline"] + width = payload.get("width") + height = payload.get("height") + if width: + pipeline.width = int(width) + if height: + pipeline.height = int(height) + + if "prompts" in payload or "prompt" in payload: + prompts = _extract_prompts(payload) + try: + io = await _apply_workflow( + pipeline, + prompts, + width=int(width) if width else None, + height=int(height) if height else None, + skip_warmup=True, + ) + except Exception as exc: + log.exception("failed to update stream workflow") + raise web.HTTPBadRequest(text=f"invalid workflow update: {exc}") from exc + session.io = io + session.prompts = prompts + if io["text"]["output"] and session.text_task is None: + if session.text_out is None and session.text_url: + session.text_out = JSONLWriter(session.text_url) + if session.text_out is not None: + session.text_task = asyncio.create_task(_text_forward_loop(app, session)) + + return web.json_response(session.to_json()) + + +async def _handle_text(request: web.Request) -> web.Response: + session: RunnerSession | None = request.app.get("session") + if session is None: + raise web.HTTPNotFound(text="no active session") + session_id = request.headers.get("Livepeer-Session-Id", "").strip() + if session_id and session_id != session.session_id: + raise web.HTTPConflict(text="runner has a different active session") + return web.json_response( + { + "session": session.session_id, + "texts": list(session.collected_text), + } + ) + + +async def _handle_healthz(_request: web.Request) -> web.Response: + return web.json_response({"ok": True, "app": APP_ID}) + + +def main() -> None: + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(levelname)s %(message)s", + ) + args = _parse_args() + if not args.workspace: + raise SystemExit("--workspace / COMFYUI_CWD is required") + + async def _on_startup(app: web.Application) -> None: + pipeline = Pipeline( + width=args.width, + height=args.height, + cwd=args.workspace, + disable_cuda_malloc=True, + gpu_only=True, + preview_method="none", + blacklist_custom_nodes=["ComfyUI-Manager"], + bootstrap_default_prompt=not args.skip_bootstrap, + ) + await pipeline.initialize() + app["pipeline"] = pipeline + app["session"] = None + + async def _release(event: Any) -> None: + await _on_session_release(app, event) + + app["registration"] = await register_runner( # Livepeer: 1 + args.orchestrator, + secret=args.orchSecret, + runner_url=args.runner_url, + app=APP_ID, + mode="persistent", + capacity=args.capacity, + price=args.price, + currency="usd", + unit="hour", + metadata='{"modalities":"workflow-driven","surfaces":["analyze","start_stream","update_stream"]}', + on_session_release=_release, + ) + log.info( + "registered app=%s runner_id=%s orchestrator=%s runner_url=%s", + APP_ID, + app["registration"].runner_id, + app["registration"].orchestrator_url, + args.runner_url, + ) + + async def _on_cleanup(app: web.Application) -> None: + await _close_session(app) + with suppress(Exception): + await app["registration"].close() # Livepeer: 3 + + app = web.Application() + app.router.add_post("/analyze", _handle_analyze) + app.router.add_post("/start_stream", _handle_start_stream) + app.router.add_post("/update_stream", _handle_update_stream) + app.router.add_get("/text", _handle_text) + app.router.add_get("/healthz", _handle_healthz) + app.on_startup.append(_on_startup) + app.on_cleanup.append(_on_cleanup) + web.run_app(app, host=args.host, port=args.port) + + +if __name__ == "__main__": + main() diff --git a/server/live_runner_client.py b/server/live_runner_client.py new file mode 100644 index 00000000..b9ca54b1 --- /dev/null +++ b/server/live_runner_client.py @@ -0,0 +1,180 @@ +#!/usr/bin/env python3 +"""Smoke client for ComfyStream live-runner: analyze → start_stream → update_stream. + +Usage: + python server/live_runner_client.py sample.mp4 --workflow path/to/workflow.json + +Livepeer integration (grep ``# Livepeer:``): + 1. reserve_session() + 2. post_json / MediaPublish through session.app_url (orch injects session headers) + 3. stop_runner_session() +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import logging +from contextlib import suppress +from pathlib import Path +from typing import Any + +import av + +from livepeer_gateway.errors import LivepeerGatewayError +from livepeer_gateway.http import get_json, post_json +from livepeer_gateway.live_runner import stop_runner_session +from livepeer_gateway.media_publish import MediaPublish +from livepeer_gateway.selection import reserve_session + +APP_ID = "comfystream" +DEFAULT_DISCOVERY = "https://ai1.eliteencoder.net:8936/discovery" +log = logging.getLogger("comfystream-client") + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="ComfyStream live-runner smoke client.") + parser.add_argument("input", help="Input video file.") + parser.add_argument("--discovery", default=DEFAULT_DISCOVERY) + parser.add_argument("--signer", default="", help="Remote signer URL for on-chain path.") + parser.add_argument("--workflow", required=True, help="ComfyUI API-format workflow JSON.") + parser.add_argument( + "--mode", + choices=("analyze", "stream", "both"), + default="analyze", + help="Which surface to exercise (default: analyze).", + ) + parser.add_argument("--max-frames", type=int, default=30) + parser.add_argument("--width", type=int, default=512) + parser.add_argument("--height", type=int, default=512) + parser.add_argument( + "--update-workflow", + default="", + help="Optional second workflow JSON for update_stream.", + ) + return parser.parse_args() + + +def _load_workflow(path: str) -> Any: + return json.loads(Path(path).read_text(encoding="utf-8")) + + +async def _publish_frames( + publish_url: str, + input_path: str, + *, + max_frames: int, +) -> None: + publisher = MediaPublish(publish_url) + try: + container = av.open(input_path) + sent = 0 + for frame in container.decode(video=0): + await publisher.write_frame(frame) + sent += 1 + if max_frames and sent >= max_frames: + break + container.close() + log.info("published %d frames to %s", sent, publish_url) + finally: + await publisher.close() + + +async def _run_analyze(args: argparse.Namespace, workflow: Any) -> None: + session = await reserve_session( # Livepeer: 1 + discovery_url=args.discovery, + app=APP_ID, + signer_url=args.signer.strip() or None, + ) + try: + async with session: + data = await post_json( # Livepeer: 2 + f"{session.app_url.rstrip('/')}/analyze", + { + "prompts": workflow, + "width": args.width, + "height": args.height, + }, + timeout=120.0, + ) + log.info("analyze started: %s", data) + await _publish_frames(data["in"], args.input, max_frames=args.max_frames) + await asyncio.sleep(2.0) + texts = await get_json(f"{session.app_url.rstrip('/')}/text", timeout=30.0) + log.info("analyze texts: %s", texts) + except LivepeerGatewayError as exc: + raise SystemExit(f"ERROR: {exc}") from exc + finally: + with suppress(Exception): + await stop_runner_session(session) # Livepeer: 3 + + +async def _run_stream(args: argparse.Namespace, workflow: Any) -> None: + session = await reserve_session( # Livepeer: 1 + discovery_url=args.discovery, + app=APP_ID, + signer_url=args.signer.strip() or None, + ) + try: + async with session: + data = await post_json( # Livepeer: 2 + f"{session.app_url.rstrip('/')}/start_stream", + { + "prompts": workflow, + "width": args.width, + "height": args.height, + }, + timeout=120.0, + ) + log.info("stream started: %s", data) + + publish = asyncio.create_task( + _publish_frames(data["in"], args.input, max_frames=args.max_frames) + ) + if args.update_workflow: + await asyncio.sleep(1.0) + update = _load_workflow(args.update_workflow) + updated = await post_json( + f"{session.app_url.rstrip('/')}/update_stream", + {"prompts": update}, + timeout=60.0, + ) + log.info("update_stream: %s", updated) + await publish + await asyncio.sleep(1.0) + except LivepeerGatewayError as exc: + raise SystemExit(f"ERROR: {exc}") from exc + finally: + with suppress(Exception): + await stop_runner_session(session) # Livepeer: 3 + log.info("stream session stopped") + + +async def _amain() -> int: + args = _parse_args() + input_path = Path(args.input).expanduser() + if not input_path.exists(): + raise SystemExit(f"input file does not exist: {input_path}") + args.input = str(input_path) + workflow = _load_workflow(args.workflow) + if args.mode in ("analyze", "both"): + await _run_analyze(args, workflow) + if args.mode in ("stream", "both"): + await _run_stream(args, workflow) + return 0 + + +def main() -> None: + logging.basicConfig( + level=logging.INFO, + format="%(asctime)s %(levelname)s %(message)s", + ) + try: + raise SystemExit(asyncio.run(_amain())) + except KeyboardInterrupt: + raise SystemExit(130) from None + + +if __name__ == "__main__": + main() diff --git a/src/comfystream/__init__.py b/src/comfystream/__init__.py index 398a9273..a3807a55 100644 --- a/src/comfystream/__init__.py +++ b/src/comfystream/__init__.py @@ -1,9 +1,12 @@ -from .client import ComfyStreamClient -from .exceptions import ComfyStreamAudioBufferError, ComfyStreamInputTimeoutError -from .pipeline import Pipeline -from .pipeline_state import PipelineState, PipelineStateManager -from .server.metrics import MetricsManager, StreamStatsManager -from .server.utils import FPSMeter, temporary_log_level +from comfy_compatibility.imports import MAIN_PY, SITE_PACKAGES, ImportContext + +with ImportContext("comfy", "comfy_extras", "comfy.vendor", order=[SITE_PACKAGES, MAIN_PY]): + from .client import ComfyStreamClient + from .exceptions import ComfyStreamAudioBufferError, ComfyStreamInputTimeoutError + from .pipeline import Pipeline + from .pipeline_state import PipelineState, PipelineStateManager + from .server.metrics import MetricsManager, StreamStatsManager + from .server.utils import FPSMeter, temporary_log_level __all__ = [ "ComfyStreamClient", diff --git a/src/comfystream/pipeline.py b/src/comfystream/pipeline.py index 9143e591..9f26a414 100644 --- a/src/comfystream/pipeline.py +++ b/src/comfystream/pipeline.py @@ -36,7 +36,7 @@ def __init__( self, width: int = 512, height: int = 512, - comfyui_inference_log_level: Optional[int] = None, + logging_level: Optional[int] = None, auto_warmup: bool = False, bootstrap_default_prompt: bool = True, **kwargs, @@ -46,12 +46,11 @@ def __init__( Args: width: Width of the video frames (default: 512) height: Height of the video frames (default: 512) - comfyui_inference_log_level: The logging level for ComfyUI inference. - Defaults to None, using the global ComfyUI log level. + logging_level: The logging level for ComfyUI (passed to Configuration). auto_warmup: Whether to run warmup automatically after prompts are set. bootstrap_default_prompt: Whether to run the default workflow once during initialization to start ComfyUI before prompts are applied. - **kwargs: Additional arguments to pass to the ComfyStreamClient + **kwargs: Additional arguments to pass to the ComfyStreamClient to configure ComfyUI """ self.client = ComfyStreamClient(**kwargs) self.width = width @@ -64,7 +63,7 @@ def __init__( self.processed_audio_buffer = np.array([], dtype=np.int16) - self._comfyui_inference_log_level = comfyui_inference_log_level + self._comfy_logging_level = logging_level self._cached_modalities: Optional[Set[str]] = None self._cached_io_capabilities: Optional[WorkflowModality] = None self.state_manager = PipelineStateManager(self.client) @@ -693,7 +692,7 @@ async def get_processed_video_frame(self) -> av.VideoFrame: return frame # Get processed output from client - async with temporary_log_level("comfy", self._comfyui_inference_log_level): + async with temporary_log_level("comfy", self._comfy_logging_level): out_tensor = await self.client.get_video_output() processed_frame = self.video_postprocess(out_tensor) @@ -726,7 +725,7 @@ async def get_processed_audio_frame(self) -> av.AudioFrame: # Process audio if needed if frame.samples > len(self.processed_audio_buffer): - async with temporary_log_level("comfy", self._comfyui_inference_log_level): + async with temporary_log_level("comfy", self._comfy_logging_level): out_tensor = await self.client.get_audio_output() self.processed_audio_buffer = np.concatenate([self.processed_audio_buffer, out_tensor]) @@ -750,7 +749,7 @@ async def get_text_output(self) -> str | None: if not self.produces_text_output(): return None - async with temporary_log_level("comfy", self._comfyui_inference_log_level): + async with temporary_log_level("comfy", self._comfy_logging_level): out_text = await self.client.get_text_output() return out_text diff --git a/src/comfystream/scripts/README.md b/src/comfystream/scripts/README.md index e593a376..42fdc543 100644 --- a/src/comfystream/scripts/README.md +++ b/src/comfystream/scripts/README.md @@ -14,12 +14,14 @@ From the repository root: -> The `--workspace` flag is optional and will default to `$COMFY_UI_WORKSPACE` or `~/comfyui`. +> The `--workspace` flag is optional and will default to `$COMFYUI_CWD` or `~/comfyui`. ### Install custom nodes + ```bash python src/comfystream/scripts/setup_nodes.py --workspace /path/to/comfyui ``` + > The optional flag `--pull-branches` can be used to ensure the latest git changes are pulled for any custom nodes defined with a `branch` in nodes.yaml #### Using a custom nodes configuration @@ -29,6 +31,7 @@ python src/comfystream/scripts/setup_nodes.py --workspace /path/to/comfyui --con > The `--config` flag accepts a filename (searches in `configs/`), relative path, or absolute path to a custom nodes configuration file ### Download models and compile tensorrt engines + ```bash python src/comfystream/scripts/setup_models.py --workspace /path/to/comfyui ``` @@ -54,7 +57,7 @@ nodes: - "tensorrt" ``` -> The `branch` property can be substituted with a SHA-256 commit hash for pinning custom node versions +> The `branch` property can be substituted with a SHA-256 commit hash for pinning custom node versions ### Models (models.yaml) @@ -86,6 +89,6 @@ workspace/ ## Environment Variables -- `COMFY_UI_WORKSPACE`: Base directory for installation +- `COMFYUI_CWD`: Base directory for installation - `PYTHONPATH`: Defaults to workspace directory - `CUSTOM_NODES_PATH`: Custom nodes directory diff --git a/src/comfystream/scripts/__init__.py b/src/comfystream/scripts/__init__.py index abcb156f..8dbb25c4 100644 --- a/src/comfystream/scripts/__init__.py +++ b/src/comfystream/scripts/__init__.py @@ -1,7 +1,3 @@ -from utils import get_config_path, load_model_config - -from . import utils -from .setup_models import run_setup_models -from .setup_nodes import run_setup_nodes - """Setup scripts for ComfyUI streaming server""" + +__all__ = [] diff --git a/src/comfystream/scripts/build_trt.py b/src/comfystream/scripts/build_trt.py index b329ff20..2d9d1d41 100644 --- a/src/comfystream/scripts/build_trt.py +++ b/src/comfystream/scripts/build_trt.py @@ -4,34 +4,101 @@ import os import sys import time +from pathlib import Path -# Reccomended running from comfystream conda environment -# in devcontainer from the workspace/ directory, or comfystream/ if you've checked out the repo -# $> conda activate comfystream -# $> python src/comfystream/scripts/build_trt.py --model /ComfyUI/models/checkpoints/SD1.5/dreamshaper-8.safetensors --out-engine /ComfyUI/output/tensorrt/static-dreamshaper8_SD15_$stat-b-1-h-512-w-512_00001_.engine +# Globals populated after workspace setup +comfy = None +detect_version_from_model = None +get_helper_from_model = None +export_onnx = None +TRTDiffusionBackbone = None -# Paths path explicitly to use the downloaded comfyUI installation on root -ROOT_DIR = "/workspace" -COMFYUI_DIR = "/workspace/ComfyUI" -timing_cache_path = "/workspace/ComfyUI/output/tensorrt/timing_cache" -if ROOT_DIR not in sys.path: - sys.path.insert(0, ROOT_DIR) -if COMFYUI_DIR not in sys.path: - sys.path.insert(0, COMFYUI_DIR) +def setup_comfy(workspace_dir: str): + """Ensure ComfyUI workspace is importable and load comfy modules. -import comfy -import comfy.model_management -from ComfyUI.custom_nodes.ComfyUI_TensorRT.models.supported_models import ( - detect_version_from_model, - get_helper_from_model, -) -from ComfyUI.custom_nodes.ComfyUI_TensorRT.onnx_utils.export import export_onnx -from ComfyUI.custom_nodes.ComfyUI_TensorRT.tensorrt_diffusion_model import TRTDiffusionBackbone + For TensorRT engine building, we use the cloned ComfyUI workspace directly + (not the pip-installed comfyui package) because custom nodes like + ComfyUI_TensorRT expect the traditional ComfyUI module structure. + """ + global \ + comfy, \ + detect_version_from_model, \ + get_helper_from_model, \ + export_onnx, \ + TRTDiffusionBackbone + + # Normalize and export the workspace so downstream imports/tools see it + workspace_dir = str(Path(workspace_dir).expanduser().resolve()) + os.environ["COMFYUI_CWD"] = workspace_dir + os.environ["COMFYUI_WORKSPACE"] = workspace_dir + + print(f"[build_trt] Using COMFYUI_CWD={workspace_dir}") + + # Ensure workspace directories have __init__.py so they are proper packages + # (not namespace packages) and take priority over pip-installed versions + workspace_path = Path(workspace_dir) + package_dirs = ["comfy_extras"] + for pkg_dir in package_dirs: + init_file = workspace_path / pkg_dir / "__init__.py" + if init_file.parent.exists() and not init_file.exists(): + init_file.touch() + print(f"[build_trt] Created {init_file}") + + # Add workspace and custom_nodes to sys.path FIRST so they take priority + custom_nodes_dir = str(workspace_path / "custom_nodes") + + # Insert at the beginning so workspace takes priority over site-packages + if workspace_dir not in sys.path: + sys.path.insert(0, workspace_dir) + if custom_nodes_dir not in sys.path: + sys.path.insert(0, custom_nodes_dir) + + # Clear any pip-installed comfy modules from sys.modules so the workspace + # versions are imported instead. The pip-installed comfyui package has + # __init__.py files which would otherwise take priority over the workspace's + # namespace packages. + modules_to_clear = ["comfy", "comfy_extras", "nodes"] + for mod_prefix in modules_to_clear: + to_delete = [key for key in sys.modules if key == mod_prefix or key.startswith(f"{mod_prefix}.")] + for key in to_delete: + del sys.modules[key] + + # Now import comfy from the workspace + import comfy as _comfy + import comfy.model_management as _cm + import comfy.sd as _sd + + comfy = _comfy + comfy.model_management = _cm + comfy.sd = _sd + + # Import TensorRT custom node modules + from ComfyUI_TensorRT.models.supported_models import ( + detect_version_from_model as _detect, + get_helper_from_model as _get_helper, + ) + from ComfyUI_TensorRT.onnx_utils.export import export_onnx as _export + from ComfyUI_TensorRT.tensorrt_diffusion_model import ( + TRTDiffusionBackbone as _TRTBackbone, + ) + + detect_version_from_model = _detect + get_helper_from_model = _get_helper + export_onnx = _export + TRTDiffusionBackbone = _TRTBackbone def parse_args(): parser = argparse.ArgumentParser(description="Build a TensorRT engine from a ComfyUI model.") + parser.add_argument( + "--workspace", + type=str, + default=os.environ.get( + "COMFYUI_CWD", os.environ.get("COMFYUI_WORKSPACE", str(Path.home() / "ComfyUI")) + ), + help="Path to the ComfyUI workspace (default: $COMFYUI_CWD, else $COMFYUI_WORKSPACE, else ~/ComfyUI)", + ) parser.add_argument( "--model", type=str, @@ -109,6 +176,7 @@ def build_trt_engine( num_video_frames: int = 14, fp8: bool = False, verbose: bool = False, + workspace_dir: str | None = None, ): """ 1) Load the model from ComfyUI by path or name @@ -135,6 +203,9 @@ def build_trt_engine( ) # 1) Load model in GPU: + if workspace_dir: + setup_comfy(workspace_dir) + comfy.model_management.unload_all_models() loaded_model = comfy.sd.load_diffusion_model(model_path, model_options={}) @@ -199,7 +270,9 @@ def build_trt_engine( # The tensorrt_diffusion_model build() signature is typically: # build(onnx_path, engine_path, timing_cache_path, opt_config, min_config, max_config) # If you have a separate 'timing_cache.trt', put it next to this script: - timing_cache_path = os.path.join(os.path.dirname(__file__), "timing_cache.trt") + timing_cache_path = os.path.join( + workspace_dir or os.path.dirname(__file__), "output", "tensorrt", "timing_cache" + ) if verbose: print(f"[INFO] Building engine -> {engine_out_path}") @@ -231,6 +304,7 @@ def build_trt_engine( def main(): args = parse_args() + setup_comfy(args.workspace) build_trt_engine( model_path=args.model, engine_out_path=args.out_engine, @@ -244,6 +318,7 @@ def main(): context_opt=args.context, fp8=args.fp8, verbose=args.verbose, + workspace_dir=args.workspace, ) diff --git a/src/comfystream/scripts/constraints.txt b/src/comfystream/scripts/constraints.txt index 54bd7cb6..3052eb5b 100644 --- a/src/comfystream/scripts/constraints.txt +++ b/src/comfystream/scripts/constraints.txt @@ -6,11 +6,14 @@ torchvision==0.23.0+cu128 torchaudio==2.8.0+cu128 tensorrt==10.12.0.36 tensorrt-cu12==10.12.0.36 +nvidia-modelopt>=0.25.0 xformers==0.0.32.post2 onnx==1.18.0 -onnxruntime>=1.22.0 -onnxruntime-gpu>=1.22.0 +onnxruntime==1.23.2 +onnxruntime-gpu==1.23.2 onnxmltools==1.14.0 cuda-python<13.0 -huggingface-hub>=0.20.0 +huggingface-hub>=0.25.0 mediapipe==0.10.21 +transformers>=4.49.0,<=4.56.0 +tokenizers<0.22 diff --git a/src/comfystream/scripts/overrides.txt b/src/comfystream/scripts/overrides.txt new file mode 100644 index 00000000..d0b3425a --- /dev/null +++ b/src/comfystream/scripts/overrides.txt @@ -0,0 +1,12 @@ +--extra-index-url https://download.pytorch.org/whl/cu128 +--extra-index-url https://pypi.nvidia.com +numpy<2.0.0 +torch==2.7.1+cu128 +torchvision==0.22.1+cu128 +torchaudio==2.7.1+cu128 +tensorrt==10.12.0.36 +tensorrt-cu12==10.12.0.36 +onnx==1.18.0 +onnxruntime==1.22.0 +onnxruntime-gpu==1.22.0 +onnxmltools==1.14.0 diff --git a/src/comfystream/scripts/setup_models.py b/src/comfystream/scripts/setup_models.py index b607bdd9..0dd94520 100644 --- a/src/comfystream/scripts/setup_models.py +++ b/src/comfystream/scripts/setup_models.py @@ -6,7 +6,8 @@ import requests import yaml from tqdm import tqdm -from utils import get_config_path, load_model_config + +from .utils import get_config_path, load_model_config try: from huggingface_hub import snapshot_download, hf_hub_download @@ -20,8 +21,10 @@ def parse_args(): parser = argparse.ArgumentParser(description="Setup ComfyUI models") parser.add_argument( "--workspace", - default=os.environ.get("COMFY_UI_WORKSPACE", os.path.expanduser("~/comfyui")), - help="ComfyUI workspace directory (default: ~/comfyui or $COMFY_UI_WORKSPACE)", + "--cwd", + dest="workspace", + default=os.environ.get("COMFYUI_CWD", os.path.expanduser("~/comfyui")), + help="ComfyUI workspace directory (default: ~/comfyui or $COMFYUI_CWD)", ) parser.add_argument('--config', default=None, @@ -206,6 +209,8 @@ def setup_models(): sys.exit(1) setup_directories(workspace_dir) - setup_model_files(workspace_dir) + setup_model_files(workspace_dir, config_path) + -setup_models() +if __name__ == "__main__": + setup_models() diff --git a/src/comfystream/scripts/setup_nodes.py b/src/comfystream/scripts/setup_nodes.py index 9bacc61a..b6b00bd8 100755 --- a/src/comfystream/scripts/setup_nodes.py +++ b/src/comfystream/scripts/setup_nodes.py @@ -5,15 +5,18 @@ from pathlib import Path import yaml -from utils import get_config_path, load_model_config + +CONSTRAINTS_PATH = Path(__file__).parent / "constraints.txt" def parse_args(): parser = argparse.ArgumentParser(description="Setup ComfyUI nodes and models") parser.add_argument( "--workspace", - default=os.environ.get("COMFY_UI_WORKSPACE", Path("~/comfyui").expanduser()), - help="ComfyUI workspace directory (default: ~/comfyui or $COMFY_UI_WORKSPACE)", + "--cwd", + dest="workspace", + default=os.environ.get("COMFYUI_CWD", Path("~/comfyui").expanduser()), + help="ComfyUI workspace directory (default: ~/comfyui or $COMFYUI_CWD)", ) parser.add_argument( "--pull-branches", @@ -30,14 +33,13 @@ def parse_args(): def setup_environment(workspace_dir): - os.environ["COMFY_UI_WORKSPACE"] = str(workspace_dir) + os.environ["COMFYUI_CWD"] = str(workspace_dir) os.environ["PYTHONPATH"] = str(workspace_dir) os.environ["CUSTOM_NODES_PATH"] = str(workspace_dir / "custom_nodes") def setup_directories(workspace_dir): """Create required directories in the workspace""" - # Create base directories workspace_dir.mkdir(parents=True, exist_ok=True) custom_nodes_dir = workspace_dir / "custom_nodes" custom_nodes_dir.mkdir(parents=True, exist_ok=True) @@ -46,9 +48,11 @@ def setup_directories(workspace_dir): def install_custom_nodes(workspace_dir, config_path=None, pull_branches=False): """Install custom nodes based on configuration""" if config_path is None: - config_path = get_config_path("nodes.yaml") + config_path = Path("configs") / "nodes.yaml" + try: - config = load_model_config(config_path) + with open(config_path, "r") as f: + config = yaml.safe_load(f) except FileNotFoundError: print(f"Error: Nodes config file not found at {config_path}") return @@ -60,14 +64,13 @@ def install_custom_nodes(workspace_dir, config_path=None, pull_branches=False): custom_nodes_path.mkdir(parents=True, exist_ok=True) os.chdir(custom_nodes_path) - # Get the absolute path to constraints.txt - constraints_path = Path(__file__).parent / "constraints.txt" - if not constraints_path.exists(): - print(f"Warning: constraints.txt not found at {constraints_path}") - constraints_path = None + failed_nodes = [] - try: - for _, node_info in config["nodes"].items(): + # Build constraints args once, used for all pip installs + constraints_args = ["-c", str(CONSTRAINTS_PATH)] if CONSTRAINTS_PATH.exists() else [] + + for _, node_info in config["nodes"].items(): + try: dir_name = node_info["url"].split("/")[-1].replace(".git", "") node_path = custom_nodes_path / dir_name @@ -98,36 +101,82 @@ def install_custom_nodes(workspace_dir, config_path=None, pull_branches=False): # Install requirements if present requirements_file = node_path / "requirements.txt" if requirements_file.exists(): - pip_cmd = [ - sys.executable, - "-m", - "pip", - "install", - "-r", - str(requirements_file), - ] - if constraints_path and constraints_path.exists(): - pip_cmd.extend(["-c", str(constraints_path)]) - subprocess.run(pip_cmd, check=True) + print(f"Installing requirements from {requirements_file}") + + # Parse requirements file to extract --extra-index-url lines + # uv doesn't support these directives in requirements files + extra_index_urls = [] + package_lines = [] + + with open(requirements_file) as f: + for line in f: + line = line.strip() + if not line or line.startswith("#"): + continue + # Check if --extra-index-url is on its own line + if line.startswith("--extra-index-url"): + url = line.split(None, 1)[1] if len(line.split(None, 1)) > 1 else "" + if url: + extra_index_urls.append(url) + # Check if --extra-index-url is inline with a package + elif "--extra-index-url" in line: + parts = line.split("--extra-index-url", 1) + package = parts[0].strip() + url = parts[1].strip() if len(parts) > 1 else "" + if package: + package_lines.append(package) + if url: + extra_index_urls.append(url) + else: + # Strip [all] extra from nvidia-modelopt to avoid + # onnxruntime-gpu version conflicts + if line.startswith("nvidia-modelopt"): + line = line.replace("[all]", "") + package_lines.append(line) + + # Create temp requirements file without --extra-index-url directives + if extra_index_urls: + temp_req = node_path / "requirements.txt.tmp" + with open(temp_req, "w") as f: + f.write("\n".join(package_lines)) + + uv_cmd = ["uv", "pip", "install"] + for url in extra_index_urls: + uv_cmd.extend(["--extra-index-url", url]) + uv_cmd.extend(["-r", str(temp_req)]) + uv_cmd.extend(constraints_args) + subprocess.run(uv_cmd, check=True) + temp_req.unlink() + else: + uv_cmd = ["uv", "pip", "install", "-r", str(requirements_file)] + uv_cmd.extend(constraints_args) + subprocess.run(uv_cmd, check=True) # Install additional dependencies if specified if "dependencies" in node_info: for dep in node_info["dependencies"]: - pip_cmd = [sys.executable, "-m", "pip", "install", dep] - if constraints_path and constraints_path.exists(): - pip_cmd.extend(["-c", str(constraints_path)]) - subprocess.run(pip_cmd, check=True) + print(f"Installing dependency: {dep}") + uv_cmd = ["uv", "pip", "install"] + uv_cmd.extend(constraints_args) + uv_cmd.append(dep) + subprocess.run(uv_cmd, check=True) + + print(f"✓ Installed {node_info['name']}") + except Exception as e: + print(f"✗ Error installing {node_info['name']}: {e}") + failed_nodes.append(node_info["name"]) + continue - print(f"Installed {node_info['name']}") - except Exception as e: - print(f"Error installing {node_info['name']} {e}") - raise e + if failed_nodes: + print(f"\nWarning: {len(failed_nodes)} node(s) failed to install:") + for name in failed_nodes: + print(f" - {name}") def setup_nodes(): args = parse_args() workspace_dir = Path(args.workspace) - + # Resolve config path if provided config_path = None if args.config: diff --git a/src/comfystream/scripts/utils.py b/src/comfystream/scripts/utils.py index 5ce4ab35..1eb982fe 100644 --- a/src/comfystream/scripts/utils.py +++ b/src/comfystream/scripts/utils.py @@ -1,24 +1,82 @@ -from pathlib import Path - import yaml - +import os +from pathlib import Path def get_config_path(filename): - """Get the absolute path to a config file""" - config_path = Path("configs") / filename - if not config_path.exists(): - print(f"Warning: Config file {filename} not found at {config_path}") - print("Available files in configs/:") - try: - for f in Path("configs").glob("*"): - print(f" - {f.name}") - except FileNotFoundError: - print(" configs/ directory not found") - raise FileNotFoundError(f"Config file {filename} not found at {config_path}") - return config_path - + """Get the absolute path to a config file with pattern matching support""" + configs_dir = Path("configs") + + if not configs_dir.exists(): + print(" configs/ directory not found") + raise FileNotFoundError("configs/ directory not found") + + # First try exact match + config_path = configs_dir / filename + if config_path.exists(): + return config_path + + # If no extension provided, try adding .yaml + if not filename.endswith('.yaml'): + config_path = configs_dir / f"{filename}.yaml" + if config_path.exists(): + return config_path + + # Try pattern matching for nodes-* files + if not filename.startswith('nodes-') and not filename == 'nodes.yaml': + pattern_path = configs_dir / f"nodes-{filename}.yaml" + if pattern_path.exists(): + return pattern_path + + # If still not found, show available files + print(f"Warning: Config file matching '{filename}' not found") + raise FileNotFoundError(f"Config file matching '{filename}' not found") def load_model_config(config_path): """Load model configuration from YAML file""" - with open(config_path, "r") as f: + with open(config_path, 'r') as f: return yaml.safe_load(f) + +def get_default_workspace(): + """Get the default workspace directory""" + return os.environ.get("COMFYUI_WORKSPACE", Path("~/comfyui").expanduser()) + +def validate_and_prompt_workspace(workspace_path, script_name="script"): + """ + Validate workspace directory exists and prompt user to create if it doesn't. + + Args: + workspace_path: Path to workspace directory (str or Path object) + script_name: Name of the calling script for better error messages + + Returns: + Path: Validated workspace directory path + + Raises: + SystemExit: If user cancels workspace creation + """ + workspace_dir = Path(workspace_path) + + # Check if workspace exists, and prompt user if it doesn't + if not workspace_dir.exists(): + print(f"Workspace directory '{workspace_dir}' does not exist.") + + # Check if this is the default workspace (user didn't specify one) + default_workspace = get_default_workspace() + if str(workspace_dir) == str(default_workspace): + print("No workspace was specified and the default workspace doesn't exist.") + + try: + response = input(f"Would you like to create '{workspace_dir}' and continue? (y/N): ").strip().lower() + if response not in ['y', 'yes']: + print(f"{script_name} cancelled.") + raise SystemExit(0) + except (KeyboardInterrupt, EOFError): + print(f"\n{script_name} cancelled.") + raise SystemExit(0) + + return workspace_dir + +def setup_workspace_environment(workspace_dir): + """Setup environment variables for workspace""" + os.environ["COMFYUI_WORKSPACE"] = str(workspace_dir) + os.environ["CUSTOM_NODES_PATH"] = str(workspace_dir / "custom_nodes") diff --git a/workflows/comfystream/analyze-stub-api.json b/workflows/comfystream/analyze-stub-api.json new file mode 100644 index 00000000..0ce11cb1 --- /dev/null +++ b/workflows/comfystream/analyze-stub-api.json @@ -0,0 +1,15 @@ +{ + "1": { + "inputs": {}, + "class_type": "LoadTensor", + "_meta": { "title": "LoadTensor" } + }, + "2": { + "inputs": { + "data": "comfystream-analyze-ok", + "remove_linebreaks": true + }, + "class_type": "SaveTextTensor", + "_meta": { "title": "SaveTextTensor" } + } +}