feat(server): make ComfyUI log level configurable

This commit introduces the `comfyui-log-level` and `comfyui-inference-log-level`
arguments to the server, allowing users to adjust the log level and reduce
excessive logging.

Author: rickstaa

server/app.py

@@ -287,7 +287,10 @@ async def on_startup(app: web.Application):
         patch_loop_datagram(app["media_ports"])
 
     app["pipeline"] = Pipeline(
-        cwd=app["workspace"], disable_cuda_malloc=True, gpu_only=True
+        cwd=app["workspace"], 
+        disable_cuda_malloc=True, 
+        gpu_only=True, 
+        comfyui_inference_log_level=app.get("comfui_inference_log_level", None),
     )
     app["pcs"] = set()
     app["video_tracks"] = {}
@@ -328,6 +331,18 @@ async def on_shutdown(app: web.Application):
         action="store_true",
         help="Include stream ID as a label in Prometheus metrics.",
     )
+    parser.add_argument(
+        "--comfyui-log-level",
+        default=None,
+        choices=logging._nameToLevel.keys(),
+        help="Set the global logging level for ComfyUI",
+    )
+    parser.add_argument(
+        "--comfyui-inference-log-level",
+        default=None,
+        choices=logging._nameToLevel.keys(),
+        help="Set the logging level for ComfyUI inference",
+    )
     args = parser.parse_args()
 
     logging.basicConfig(
@@ -377,4 +392,11 @@ def force_print(*args, **kwargs):
         print(*args, **kwargs, flush=True)
         sys.stdout.flush()
 
+    # Allow overriding of ComyfUI log levels.
+    if args.comfyui_log_level:
+        log_level = logging._nameToLevel.get(args.comfyui_log_level.upper())
+        logging.getLogger("comfy").setLevel(log_level)
+    if args.comfyui_inference_log_level:
+        app["comfui_inference_log_level"] = args.comfyui_inference_log_level
+
     web.run_app(app, host=args.host, port=int(args.port), print=force_print)

server/pipeline.py

@@ -5,19 +5,28 @@
 
 from typing import Any, Dict, Union, List
 from comfystream.client import ComfyStreamClient
+from utils import temporary_log_level
 
 WARMUP_RUNS = 5
 
 
 class Pipeline:
-    def __init__(self, **kwargs):
+    def __init__(self, comfyui_inference_log_level: int = None, **kwargs):
+        """Initialize the pipeline with the given configuration.
+        Args:
+            comfyui_inference_log_level: The logging level for ComfyUI inference.
+                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.video_incoming_frames = asyncio.Queue()
         self.audio_incoming_frames = asyncio.Queue()
 
         self.processed_audio_buffer = np.array([], dtype=np.int16)
 
+        self._comfyui_inference_log_level = comfyui_inference_log_level
+
     async def warm_video(self):
         dummy_frame = av.VideoFrame()
         dummy_frame.side_data.input = torch.randn(1, 512, 512, 3)
@@ -76,7 +85,8 @@ def audio_postprocess(self, output: Union[torch.Tensor, np.ndarray]) -> av.Audio
 
     async def get_processed_video_frame(self):
         # TODO: make it generic to support purely generative video cases
-        out_tensor = await self.client.get_video_output()
+        async with temporary_log_level("comfy", self._comfyui_inference_log_level):
+            out_tensor = await self.client.get_video_output()
         frame = await self.video_incoming_frames.get()
         while frame.side_data.skipped:
             frame = await self.video_incoming_frames.get()
@@ -91,7 +101,8 @@ async def get_processed_audio_frame(self):
         # TODO: make it generic to support purely generative audio cases and also add frame skipping
         frame = await self.audio_incoming_frames.get()
         if frame.samples > len(self.processed_audio_buffer):
-            out_tensor = await self.client.get_audio_output()
+            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:]
@@ -109,4 +120,4 @@ async def get_nodes_info(self) -> Dict[str, Any]:
         return nodes_info
 
     async def cleanup(self):
-        await self.client.cleanup()
+        await self.client.cleanup()