diff --git a/configs/build_targets.yaml b/configs/build_targets.yaml new file mode 100644 index 00000000..6836ee95 --- /dev/null +++ b/configs/build_targets.yaml @@ -0,0 +1,17 @@ +# build_targets.yaml +# Maps node names to their build scripts and config files. +# Add new nodes as needed. + +streamdiffusion: + script: /workspace/ComfyUI/custom_nodes/ComfyUI-StreamDiffusion/scripts/build_tensorrt_engines.py + configs: + - /workspace/ComfyUI/custom_nodes/ComfyUI-StreamDiffusion/configs/sd15_singlecontrol.yaml + - /workspace/ComfyUI/custom_nodes/ComfyUI-StreamDiffusion/configs/sdturbo_multicontrol.yaml + folder: /workspace/ComfyUI/custom_nodes/ComfyUI-StreamDiffusion/scripts + +# Example for another node: +# depthanything: +# script: /workspace/ComfyUI/custom_nodes/ComfyUI-DepthAnything/scripts/build_depthanything_engine.py +# configs: +# - /workspace/ComfyUI/custom_nodes/ComfyUI-DepthAnything/configs/depthanything.yaml +# folder: /workspace/ComfyUI/custom_nodes/ComfyUI-DepthAnything/scripts diff --git a/docker/Dockerfile.base b/docker/Dockerfile.base index d2ce3fbb..c65f849d 100644 --- a/docker/Dockerfile.base +++ b/docker/Dockerfile.base @@ -56,6 +56,9 @@ RUN apt-get remove --purge -y libcudnn9-cuda-12 libcudnn9-dev-cuda-12 || true && # to ensure numpy 2.0 is not installed automatically by another package RUN conda run -n comfystream --no-capture-output pip install "numpy<2.0.0" +# Ensure modelopt pulls a compatible transformers version for HF support +RUN conda run -n comfystream --no-capture-output pip install 'nvidia-modelopt[hf]' + # Install cuDNN 9.8 via conda to match base system version # Caution: Mixed versions installed in environment (system/python) can cause CUDNN_STATUS_SUBLIBRARY_VERSION_MISMATCH errors RUN conda install -n comfystream -y -c nvidia -c conda-forge cudnn=9.8 cuda-version=12.8 diff --git a/docker/entrypoint.sh b/docker/entrypoint.sh index c37e8d76..8852d5c8 100755 --- a/docker/entrypoint.sh +++ b/docker/entrypoint.sh @@ -138,17 +138,40 @@ if [ "$1" = "--build-engines" ]; then echo "Engines for FasterLivePortrait already exists, skipping..." fi - # Build Engine for StreamDiffusion using trt script and config - ENGINE_SCRIPT="/workspace/ComfyUI/custom_nodes/ComfyUI-StreamDiffusion/scripts/build_tensorrt_engines.py" - CONFIGS=( - "/workspace/ComfyUI/custom_nodes/ComfyUI-StreamDiffusion/configs/sd15_singlecontrol.yaml" - "/workspace/ComfyUI/custom_nodes/ComfyUI-StreamDiffusion/configs/sdturbo_multicontrol.yaml" - ) - cd /workspace/ComfyUI/custom_nodes/ComfyUI-StreamDiffusion/scripts - for ENGINE_CONFIG in "${CONFIGS[@]}"; do - echo "Building StreamDiffusion TensorRT engines using config: $ENGINE_CONFIG" - python "$ENGINE_SCRIPT" --config "$ENGINE_CONFIG" - done + # Build Engine for StreamDiffusion using build_targets.yaml (dynamic, robust) + BUILD_TARGETS_FILE="/workspace/comfystream/configs/build_targets.yaml" + if [ ! -f "$BUILD_TARGETS_FILE" ]; then + echo "build_targets.yaml not found at $BUILD_TARGETS_FILE. Skipping StreamDiffusion engine builds." + else + python3 - <<'EOF' +import os +import yaml +import subprocess + +with open("/workspace/comfystream/configs/build_targets.yaml", "r") as f: + targets = yaml.safe_load(f) + +info = targets.get("streamdiffusion") +if info: + folder = info.get("folder") + script = info.get("script") + configs = info.get("configs", []) + if folder and os.path.isdir(folder) and script and os.path.isfile(script): + for config in configs: + if os.path.isfile(config): + print(f"Building streamdiffusion engine with config: {config}") + try: + subprocess.run(["python", script, "--config", config], check=True) + except subprocess.CalledProcessError as e: + print(f"Error building engine for config {config}: {e}") + else: + print(f"Warning: Config {config} for streamdiffusion not found, skipping...") + else: + print(f"Skipping streamdiffusion: required folder or script not found.") +else: + print("streamdiffusion not found in build_targets.yaml, skipping...") +EOF + fi shift fi @@ -218,4 +241,6 @@ if [ "$START_COMFYUI" = true ] || [ "$START_API" = true ] || [ "$START_UI" = tru tail -f /var/log/supervisord.log fi + + exec "$@" diff --git a/src/comfystream/scripts/constraints.txt b/src/comfystream/scripts/constraints.txt index b36010e6..54bd7cb6 100644 --- a/src/comfystream/scripts/constraints.txt +++ b/src/comfystream/scripts/constraints.txt @@ -8,8 +8,9 @@ tensorrt==10.12.0.36 tensorrt-cu12==10.12.0.36 xformers==0.0.32.post2 onnx==1.18.0 -onnxruntime==1.22.0 -onnxruntime-gpu==1.22.0 +onnxruntime>=1.22.0 +onnxruntime-gpu>=1.22.0 onnxmltools==1.14.0 cuda-python<13.0 huggingface-hub>=0.20.0 +mediapipe==0.10.21