Merge branch 'main' into add_comfyui_log_level_arg

Author: rickstaa

.devcontainer/post-create.sh

@@ -7,10 +7,16 @@ cd /workspace/comfystream
 echo -e "\e[32mInstalling Comfystream in editable mode...\e[0m"
 /workspace/miniconda3/envs/comfystream/bin/python3 -m pip install -e . --root-user-action=ignore > /dev/null
 
+# Install npm packages if needed
+if [ ! -d "/workspace/comfystream/ui/node_modules" ]; then
+    echo -e "\e[32mInstalling npm packages for Comfystream UI...\e[0m"
+    cd /workspace/comfystream/ui
+    npm install
+fi
+
 if [ ! -d "/workspace/comfystream/nodes/web/static" ]; then
     echo -e "\e[32mBuilding web assets...\e[0m"
     cd /workspace/comfystream/ui
-    npm install
     npm run build
 fi
 

.github/workflows/comfyui-base.yaml

@@ -40,7 +40,7 @@ jobs:
           ref: ${{ github.event.pull_request.head.sha }}
 
       - name: Login to DockerHub
-        uses: docker/login-action@v2
+        uses: docker/login-action@v3
         with:
           username: ${{ secrets.CI_DOCKERHUB_USERNAME }}
           password: ${{ secrets.CI_DOCKERHUB_TOKEN }}

.github/workflows/docker.yaml

@@ -30,7 +30,7 @@ jobs:
           ref: ${{ github.event.pull_request.head.sha }}
 
       - name: Login to DockerHub
-        uses: docker/login-action@v2
+        uses: docker/login-action@v3
         with:
           username: ${{ secrets.CI_DOCKERHUB_USERNAME }}
           password: ${{ secrets.CI_DOCKERHUB_TOKEN }}

docker/entrypoint.sh

@@ -62,11 +62,6 @@ if [ "$1" = "--build-engines" ]; then
   shift
 fi
 
-# Install npm packages if needed
-cd /workspace/comfystream/ui
-if [ ! -d "node_modules" ]; then
-  npm install --legacy-peer-deps
-fi
 
 if [ "$1" = "--server" ]; then
   /usr/bin/supervisord -c /etc/supervisor/supervisord.conf

pyproject.toml

@@ -10,10 +10,11 @@ license = { file = "LICENSE" }
 dependencies = [
     "asyncio",
     "comfyui @ git+https://github.com/hiddenswitch/ComfyUI.git@ce3583ad42c024b8f060d0002cbe20c265da6dc8",
-    "toml",
     "aiortc",
     "aiohttp",
+    "toml",
     "twilio",
+    "prometheus_client",
 ]
 
 [project.optional-dependencies]

requirements.txt

@@ -3,4 +3,5 @@ comfyui @ git+https://github.com/hiddenswitch/ComfyUI.git@ce3583ad42c024b8f060d0
 aiortc
 aiohttp
 toml
-twilio
+twilio
+prometheus_client

server/app.py

@@ -24,7 +24,9 @@
 from aiortc.rtcrtpsender import RTCRtpSender
 from pipeline import Pipeline
 from twilio.rest import Client
-from utils import FPSMeter, StreamStats, add_prefix_to_app_routes, patch_loop_datagram
+from utils import patch_loop_datagram, add_prefix_to_app_routes, FPSMeter
+from metrics import MetricsManager, StreamStatsManager
+import time
 
 logger = logging.getLogger(__name__)
 logging.getLogger("aiortc.rtcrtpsender").setLevel(logging.WARNING)
@@ -56,7 +58,9 @@ def __init__(self, track: MediaStreamTrack, pipeline: Pipeline):
         super().__init__()
         self.track = track
         self.pipeline = pipeline
-        self.fps_meter = FPSMeter()
+        self.fps_meter = FPSMeter(
+            metrics_manager=app["metrics_manager"], track_id=track.id
+        )
 
         asyncio.create_task(self.collect_frames())
 
@@ -327,6 +331,18 @@ async def on_shutdown(app: web.Application):
         choices=logging._nameToLevel.keys(),
         help="Set the logging level for ComfyUI inference",
     )
+    parser.add_argument(
+        "--monitor",
+        default=False,
+        action="store_true",
+        help="Start a Prometheus metrics endpoint for monitoring.",
+    )
+    parser.add_argument(
+        "--stream-id-label",
+        default=False,
+        action="store_true",
+        help="Include stream ID as a label in Prometheus metrics.",
+    )
     args = parser.parse_args()
 
     logging.basicConfig(
@@ -350,11 +366,23 @@ async def on_shutdown(app: web.Application):
     app.router.add_post("/prompt", set_prompt)
 
     # Add routes for getting stream statistics.
-    stream_stats = StreamStats(app)
-    app.router.add_get("/streams/stats", stream_stats.collect_all_stream_metrics)
+    stream_stats_manager = StreamStatsManager(app)
     app.router.add_get(
-        "/stream/{stream_id}/stats", stream_stats.collect_stream_metrics_by_id
+        "/streams/stats", stream_stats_manager.collect_all_stream_metrics
     )
+    app.router.add_get(
+        "/stream/{stream_id}/stats", stream_stats_manager.collect_stream_metrics_by_id
+    )
+
+    # Add Prometheus metrics endpoint.
+    app["metrics_manager"] = MetricsManager(include_stream_id=args.stream_id_label)
+    if args.monitor:
+        app["metrics_manager"].enable()
+        logger.info(
+            f"Monitoring enabled - Prometheus metrics available at: "
+            f"http://{args.host}:{args.port}/metrics"
+        )
+        app.router.add_get("/metrics", app["metrics_manager"].metrics_handler)
 
     # Add hosted platform route prefix.
     # NOTE: This ensures that the local and hosted experiences have consistent routes.

server/metrics/__init__.py

@@ -0,0 +1,2 @@
+from .prometheus_metrics import MetricsManager
+from .stream_stats import StreamStatsManager

server/metrics/prometheus_metrics.py

@@ -0,0 +1,44 @@
+"""Prometheus metrics utilities."""
+
+from prometheus_client import Gauge, generate_latest
+from aiohttp import web
+from typing import Optional
+
+
+class MetricsManager:
+    """Manages Prometheus metrics collection."""
+
+    def __init__(self, include_stream_id: bool = False):
+        """Initializes the MetricsManager class.
+
+        Args:
+            include_stream_id: Whether to include the stream ID as a label in the metrics.
+        """
+        self._enabled = False
+        self._include_stream_id = include_stream_id
+
+        base_labels = ["stream_id"] if include_stream_id else []
+        self._fps_gauge = Gauge(
+            "stream_fps", "Frames per second of the stream", base_labels
+        )
+
+    def enable(self):
+        """Enable Prometheus metrics collection."""
+        self._enabled = True
+
+    def update_fps_metrics(self, fps: float, stream_id: Optional[str] = None):
+        """Update Prometheus metrics for a given stream.
+
+        Args:
+            fps: The current frames per second.
+            stream_id: The ID of the stream.
+        """
+        if self._enabled:
+            if self._include_stream_id:
+                self._fps_gauge.labels(stream_id=stream_id or "").set(fps)
+            else:
+                self._fps_gauge.set(fps)
+
+    async def metrics_handler(self, _):
+        """Handle Prometheus metrics endpoint."""
+        return web.Response(body=generate_latest(), content_type="text/plain")

server/metrics/stream_stats.py

@@ -0,0 +1,76 @@
+"""Handles real-time video stream statistics (non-Prometheus, JSON API)."""
+
+from typing import Any, Dict
+import json
+from aiohttp import web
+from aiortc import MediaStreamTrack
+
+
+class StreamStatsManager:
+    """Handles real-time video stream statistics collection."""
+
+    def __init__(self, app: web.Application):
+        """Initializes the StreamMetrics class.
+
+        Args:
+            app: The web application instance storing stream tracks.
+        """
+        self._app = app
+
+    async def collect_video_metrics(
+        self, video_track: MediaStreamTrack
+    ) -> Dict[str, Any]:
+        """Collects real-time statistics for a video track.
+
+        Args:
+            video_track: The video stream track instance.
+
+        Returns:
+            A dictionary containing FPS-related statistics.
+        """
+        return {
+            "timestamp": await video_track.fps_meter.last_fps_calculation_time,
+            "fps": await video_track.fps_meter.fps,
+            "minute_avg_fps": await video_track.fps_meter.average_fps,
+            "minute_fps_array": await video_track.fps_meter.fps_measurements,
+        }
+
+    async def collect_all_stream_metrics(self, _) -> web.Response:
+        """Retrieves real-time metrics for all active video streams.
+
+        Returns:
+            A JSON response containing FPS statistics for all streams.
+        """
+        video_tracks = self._app.get("video_tracks", {})
+        all_stats = {
+            stream_id: await self.collect_video_metrics(track)
+            for stream_id, track in video_tracks.items()
+        }
+
+        return web.Response(
+            content_type="application/json",
+            text=json.dumps(all_stats),
+        )
+
+    async def collect_stream_metrics_by_id(self, request: web.Request) -> web.Response:
+        """Retrieves real-time metrics for a specific video stream by ID.
+
+        Args:
+            request: The HTTP request containing the stream ID.
+
+        Returns:
+            A JSON response with stream metrics or an error message.
+        """
+        stream_id = request.match_info.get("stream_id")
+        video_tracks = self._app.get("video_tracks", {})
+        video_track = video_tracks.get(stream_id)
+
+        if video_track:
+            stats = await self.collect_video_metrics(video_track)
+        else:
+            stats = {"error": "Stream not found"}
+
+        return web.Response(
+            content_type="application/json",
+            text=json.dumps(stats),
+        )

server/pipeline.py

@@ -19,9 +19,7 @@ def __init__(self, comfyui_inference_log_level: int = None, **kwargs):
                 Defaults to None, using the global ComfyUI log level.
             **kwargs: Additional arguments to pass to the ComfyStreamClient
         """
-        self.client = ComfyStreamClient(
-            **kwargs, max_workers=5
-        )  # TODO: hardcoded max workers, should it be configurable?
+        self.client = ComfyStreamClient(**kwargs, max_workers=5) # TODO: hardcoded max workers, should it be configurable?
 
         self.video_incoming_frames = asyncio.Queue()
         self.audio_incoming_frames = asyncio.Queue()
@@ -40,9 +38,7 @@ async def warm_video(self):
 
     async def warm_audio(self):
         dummy_frame = av.AudioFrame()
-        dummy_frame.side_data.input = np.random.randint(
-            -32768, 32767, int(48000 * 0.5), dtype=np.int16
-        )  # TODO: adds a lot of delay if it doesn't match the buffer size, is warmup needed?
+        dummy_frame.side_data.input = np.random.randint(-32768, 32767, int(48000 * 0.5), dtype=np.int16) # TODO: adds a lot of delay if it doesn't match the buffer size, is warmup needed?
         dummy_frame.sample_rate = 48000
 
         for _ in range(WARMUP_RUNS):
@@ -55,9 +51,7 @@ async def set_prompts(self, prompts: Union[Dict[Any, Any], List[Dict[Any, Any]]]
         else:
             await self.client.set_prompts([prompts])
 
-    async def update_prompts(
-        self, prompts: Union[Dict[Any, Any], List[Dict[Any, Any]]]
-    ):
+    async def update_prompts(self, prompts: Union[Dict[Any, Any], List[Dict[Any, Any]]]):
         if isinstance(prompts, list):
             await self.client.update_prompts(prompts)
         else:
@@ -82,21 +76,12 @@ def video_preprocess(self, frame: av.VideoFrame) -> Union[torch.Tensor, np.ndarr
     def audio_preprocess(self, frame: av.AudioFrame) -> Union[torch.Tensor, np.ndarray]:
         return frame.to_ndarray().ravel().reshape(-1, 2).mean(axis=1).astype(np.int16)
 
-    def video_postprocess(
-        self, output: Union[torch.Tensor, np.ndarray]
-    ) -> av.VideoFrame:
+    def video_postprocess(self, output: Union[torch.Tensor, np.ndarray]) -> av.VideoFrame:
         return av.VideoFrame.from_ndarray(
-            (output * 255.0)
-            .clamp(0, 255)
-            .to(dtype=torch.uint8)
-            .squeeze(0)
-            .cpu()
-            .numpy()
+            (output * 255.0).clamp(0, 255).to(dtype=torch.uint8).squeeze(0).cpu().numpy()
         )
 
-    def audio_postprocess(
-        self, output: Union[torch.Tensor, np.ndarray]
-    ) -> av.AudioFrame:
+    def audio_postprocess(self, output: Union[torch.Tensor, np.ndarray]) -> av.AudioFrame:
         return av.AudioFrame.from_ndarray(np.repeat(output, 2).reshape(1, -1))
 
     async def get_processed_video_frame(self):
@@ -107,7 +92,7 @@ async def get_processed_video_frame(self):
         while frame.side_data.skipped:
             frame = await self.video_incoming_frames.get()
 
-        processed_frame = self.video_postprocess(out_tensor)
+        processed_frame  = self.video_postprocess(out_tensor)
         processed_frame.pts = frame.pts
         processed_frame.time_base = frame.time_base
 
@@ -119,17 +104,15 @@ async def get_processed_audio_frame(self):
         if frame.samples > len(self.processed_audio_buffer):
             async with temporary_log_level("comfy", self._comfyui_inference_log_level):
                 out_tensor = await self.client.get_audio_output()
-            self.processed_audio_buffer = np.concatenate(
-                [self.processed_audio_buffer, out_tensor]
-            )
-        out_data = self.processed_audio_buffer[: frame.samples]
-        self.processed_audio_buffer = self.processed_audio_buffer[frame.samples :]
+            self.processed_audio_buffer = np.concatenate([self.processed_audio_buffer, out_tensor])
+        out_data = self.processed_audio_buffer[:frame.samples]
+        self.processed_audio_buffer = self.processed_audio_buffer[frame.samples:]
 
         processed_frame = self.audio_postprocess(out_data)
         processed_frame.pts = frame.pts
         processed_frame.time_base = frame.time_base
         processed_frame.sample_rate = frame.sample_rate
-
+        
         return processed_frame
 
     async def get_nodes_info(self) -> Dict[str, Any]: