Merge branch 'main' into comfystream-mgmt-api

Author: ad-astra-video

.devcontainer/devcontainer.json

@@ -17,7 +17,7 @@
 				"python.defaultInterpreterPath": "/workspace/miniconda3/envs/comfystream/bin/python",
 				"python.venvPath": "/workspace/miniconda3/envs",
 				"python.terminal.activateEnvInCurrentTerminal": false,
-				"python.terminal.activateEnvironment": true,
+				"python.terminal.activateEnvironment": false,
 				"terminal.integrated.shellIntegration.enabled": true
 			},
 			"extensions": [

.gitignore

@@ -22,3 +22,6 @@ launch.json
 nodes/web/static/*
 .cursor/
 !nodes/web/static/.gitkeep
+
+# opencv files
+opencv

configs/models.yaml

@@ -16,11 +16,15 @@ models:
       - url: "https://huggingface.co/aaronb/dreamshaper-8-dmd-1kstep/raw/main/config.json"
         path: "unet/dreamshaper-8-dmd-1kstep.json"
 
-  # Depth Anything ONNX model
+  # Depth Anything V2 ONNX models
   depthanything-onnx:
     name: "DepthAnything ONNX"
     url: "https://huggingface.co/yuvraj108c/Depth-Anything-2-Onnx/resolve/main/depth_anything_v2_vitb.onnx?download=true"
     path: "tensorrt/depth-anything/depth_anything_vitl14.onnx"
+  depth-anything-v2-large-onnx:
+    name: "DepthAnything V2 Large ONNX"
+    url: "https://huggingface.co/yuvraj108c/Depth-Anything-2-Onnx/resolve/main/depth_anything_v2_vitl.onnx?download=true"
+    path: "tensorrt/depth-anything/depth_anything_v2_vitl.onnx"
 
   # TAESD models
   taesd:
@@ -68,4 +72,4 @@ models:
     name: "ClipTextModel"
     url: "https://huggingface.co/Lykon/dreamshaper-8/resolve/main/text_encoder/model.fp16.safetensors"
     path: "text_encoders/CLIPText/model.fp16.safetensors"
-    type: "text_encoder"
+    type: "text_encoder"

configs/nodes.yaml

@@ -11,7 +11,8 @@ nodes:
 
   comfyui-depthanything-tensorrt:
     name: "ComfyUI DepthAnything TensorRT"
-    url: "https://github.com/yuvraj108c/ComfyUI-Depth-Anything-Tensorrt"
+    url: "https://github.com/rickstaa/ComfyUI-Depth-Anything-Tensorrt"
+    branch: "feature/add-export-trt-args"
     type: "tensorrt"
 
   # Ryan's nodes
@@ -74,4 +75,4 @@ nodes:
     name: "ComfyUI Stream Pack"
     url: "https://github.com/livepeer/ComfyUI-Stream-Pack"
     branch: "main"
-    type: "utility"
+    type: "utility"

docker/Dockerfile

@@ -23,7 +23,6 @@ RUN	bash -c "source $NVM_DIR/nvm.sh && \
 ENV	NODE_PATH="$NVM_DIR/v$NODE_VERSION/lib/node_modules" \
 	PATH="$NVM_DIR/versions/node/v$NODE_VERSION/bin:$PATH"
 
-RUN	conda init bash
 
 # Create the supervisor configuration file for ComfyUI and ComfyStream
 COPY	docker/supervisord.conf /etc/supervisor/conf.d/supervisord.conf

docker/Dockerfile.base

@@ -27,19 +27,15 @@ RUN mkdir -p /workspace/comfystream && \
     wget "https://repo.anaconda.com/miniconda/Miniconda3-${CONDA_VERSION}-Linux-x86_64.sh" -O /tmp/miniconda.sh && \
     bash /tmp/miniconda.sh -b -p /workspace/miniconda3 && \
     eval "$(/workspace/miniconda3/bin/conda shell.bash hook)" && \
-    conda create -n comfystream python="${PYTHON_VERSION}" -y && \
-    rm /tmp/miniconda.sh && \
-    conda run -n comfystream --no-capture-output pip install aiortc aiohttp requests tqdm pyyaml --root-user-action=ignore
+    conda create -n comfystream python="${PYTHON_VERSION}" ffmpeg=6 -c conda-forge -y && \
+    rm /tmp/miniconda.sh && echo 'export LD_LIBRARY_PATH=/workspace/miniconda3/envs/comfystream/lib:$LD_LIBRARY_PATH' >> ~/.bashrc
 
 # Clone ComfyUI
-ADD --link https://github.com/comfyanonymous/ComfyUI.git /workspace/ComfyUI
+RUN git clone https://github.com/comfyanonymous/ComfyUI.git /workspace/ComfyUI
 
 # Copy only files needed for setup
-COPY --link ./src/comfystream/scripts /workspace/comfystream/src/comfystream/scripts
-COPY --link ./configs /workspace/comfystream/configs
-
-# Run setup_nodes (cached unless setup_nodes.py or nodes/ changes)
-RUN conda run -n comfystream --no-capture-output --cwd /workspace/comfystream python src/comfystream/scripts/setup_nodes.py --workspace /workspace/ComfyUI
+COPY ./src/comfystream/scripts /workspace/comfystream/src/comfystream/scripts
+COPY ./configs /workspace/comfystream/configs
 
 # Copy ComfyStream files into ComfyUI
 COPY . /workspace/comfystream
@@ -52,14 +48,24 @@ COPY ./test/example-512x512.png /workspace/ComfyUI/input
 RUN conda run -n comfystream --no-capture-output --cwd /workspace/ComfyUI pip install -r requirements.txt --root-user-action=ignore
 
 # Install ComfyStream requirements
-RUN conda run -n comfystream --no-capture-output --cwd /workspace/comfystream pip install -r requirements.txt --root-user-action=ignore
-RUN conda run -n comfystream --no-capture-output --cwd /workspace/comfystream pip install . --root-user-action=ignore
 RUN ln -s /workspace/comfystream /workspace/ComfyUI/custom_nodes/comfystream
+RUN conda run -n comfystream --no-capture-output --cwd /workspace/comfystream pip install -e . --root-user-action=ignore
 RUN conda run -n comfystream --no-capture-output --cwd /workspace/comfystream python install.py --workspace /workspace/ComfyUI
+
+# Run setup_nodes
+RUN conda run -n comfystream --no-capture-output --cwd /workspace/comfystream python src/comfystream/scripts/setup_nodes.py --workspace /workspace/ComfyUI
+
+# Install additional dependencies
 RUN conda run -n comfystream --no-capture-output pip install --upgrade tensorrt-cu12-bindings tensorrt-cu12-libs --root-user-action=ignore
 
+# Setup opencv with CUDA support
+RUN conda run -n comfystream --no-capture-output bash /workspace/comfystream/docker/entrypoint.sh --opencv-cuda
+
 # Configure no environment activation by default
 RUN conda config --set auto_activate_base false && \
     conda init bash
 
+# Set comfystream environment as default
+RUN echo "conda activate comfystream" >> ~/.bashrc
+
 WORKDIR /workspace/comfystream

docker/entrypoint.sh

@@ -23,6 +23,7 @@ show_help() {
   echo "Options:"
   echo "  --download-models       Download default models"
   echo "  --build-engines         Build TensorRT engines for default models"
+  echo "  --opencv-cuda           Setup OpenCV with CUDA support"
   echo "  --server                Start the Comfystream server, UI and ComfyUI"
   echo "  --help                  Show this help message"
   echo ""
@@ -49,7 +50,7 @@ if [ "$1" = "--build-engines" ]; then
   # Build Static Engine for Dreamshaper
   python src/comfystream/scripts/build_trt.py --model /workspace/ComfyUI/models/unet/dreamshaper-8-dmd-1kstep.safetensors --out-engine /workspace/ComfyUI/output/tensorrt/static-dreamshaper8_SD15_\$stat-b-1-h-512-w-512_00001_.engine
 
-  # Build Engine for DepthAnything2
+  # Build Engine for Depth Anything V2
   if [ ! -f "$DEPTH_ANYTHING_DIR/depth_anything_vitl14-fp16.engine" ]; then
     if [ ! -d "$DEPTH_ANYTHING_DIR" ]; then
       mkdir -p "$DEPTH_ANYTHING_DIR"
@@ -59,9 +60,66 @@ if [ "$1" = "--build-engines" ]; then
   else
     echo "Engine for DepthAnything2 already exists, skipping..."
   fi
+
+  # Build Engine for Depth Anything2 (large)
+  if [ ! -f "$DEPTH_ANYTHING_DIR/depth_anything_v2_vitl-fp16.engine" ]; then
+    cd "$DEPTH_ANYTHING_DIR"
+    python /workspace/ComfyUI/custom_nodes/ComfyUI-Depth-Anything-Tensorrt/export_trt.py --trt-path "${DEPTH_ANYTHING_DIR}/depth_anything_v2_vitl-fp16.engine" --onnx-path "${DEPTH_ANYTHING_DIR}/depth_anything_v2_vitl.onnx"
+  else
+    echo "Engine for DepthAnything2 (large) already exists, skipping..."
+  fi
   shift
 fi
 
+if [ "$1" = "--opencv-cuda" ]; then
+  cd /workspace/comfystream
+  conda activate comfystream
+  
+  # Check if OpenCV CUDA build already exists
+  if [ ! -f "/workspace/comfystream/opencv-cuda-release.tar.gz" ]; then
+    # Download and extract OpenCV CUDA build
+    DOWNLOAD_NAME="opencv-cuda-release.tar.gz"
+    wget -q -O "$DOWNLOAD_NAME" https://github.com/JJassonn69/ComfyUI-Stream-Pack/releases/download/v1.0/opencv-cuda-release.tar.gz
+    tar -xzf "$DOWNLOAD_NAME" -C /workspace/comfystream/
+    rm "$DOWNLOAD_NAME"
+  else
+    echo "OpenCV CUDA build already exists, skipping download."
+  fi
+
+  # Install required libraries
+  apt-get update && apt-get install -y \
+    libgflags-dev \
+    libgoogle-glog-dev \
+    libjpeg-dev \
+    libavcodec-dev \
+    libavformat-dev \
+    libavutil-dev \
+    libswscale-dev
+
+  # Remove existing cv2 package
+  SITE_PACKAGES_DIR="/workspace/miniconda3/envs/comfystream/lib/python3.11/site-packages"
+  rm -rf "${SITE_PACKAGES_DIR}/cv2"*
+
+  # Copy new cv2 package
+  cp -r /workspace/comfystream/cv2 "${SITE_PACKAGES_DIR}/"
+
+  # Handle library dependencies
+  CONDA_ENV_LIB="/workspace/miniconda3/envs/comfystream/lib"
+  
+  # Remove existing libstdc++ and copy system one
+  rm -f "${CONDA_ENV_LIB}/libstdc++.so"*
+  cp /usr/lib/x86_64-linux-gnu/libstdc++.so* "${CONDA_ENV_LIB}/"
+
+  # Copy OpenCV libraries
+  cp /workspace/comfystream/opencv/build/lib/libopencv_* /usr/lib/x86_64-linux-gnu/
+
+  # remove the opencv-contrib and cv2 folders
+  rm -rf /workspace/comfystream/opencv_contrib
+  rm -rf /workspace/comfystream/cv2
+
+  echo "OpenCV CUDA installation completed"
+  shift
+fi
 
 if [ "$1" = "--server" ]; then
   /usr/bin/supervisord -c /etc/supervisor/supervisord.conf

server/app.py

@@ -345,7 +345,11 @@ 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, preview_method='none'
+        cwd=app["workspace"], 
+        disable_cuda_malloc=True, 
+        gpu_only=True, 
+        preview_method='none',
+        comfyui_inference_log_level=app.get("comfui_inference_log_level", None),
     )
     app["pcs"] = set()
     app["video_tracks"] = {}
@@ -386,6 +390,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(
@@ -439,5 +455,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,18 +5,27 @@
 
 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)
         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)
@@ -75,7 +84,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()
@@ -90,7 +100,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:]
@@ -108,4 +119,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()

server/utils/__init__.py

@@ -1,2 +1,2 @@
-from .utils import patch_loop_datagram, add_prefix_to_app_routes
+from .utils import patch_loop_datagram, add_prefix_to_app_routes, temporary_log_level
 from .fps_meter import FPSMeter