diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json index 3b45dc6a..1c184f5c 100644 --- a/.devcontainer/devcontainer.json +++ b/.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": [ diff --git a/.editorconfig b/.editorconfig index e53f3c8d..c39bef28 100644 --- a/.editorconfig +++ b/.editorconfig @@ -9,6 +9,19 @@ charset = utf-8 trim_trailing_whitespace = true insert_final_newline = true +[.gitignore] +insert_final_newline = unset + +[*.py] +indent_size = 4 + +[workflows/comfy*/*.json] +insert_final_newline = unset + [Dockerfile*] -indent_style = tab -indent_size = 8 +indent_size = 4 + +[*.{md,txt,mkdn}] +indent_size = 4 +indent_size = unset +trim_trailing_whitespace = false diff --git a/.github/codeql-config.yaml b/.github/codeql-config.yaml new file mode 100644 index 00000000..53a66416 --- /dev/null +++ b/.github/codeql-config.yaml @@ -0,0 +1,2 @@ +paths-ignore: + - "docker/" diff --git a/.github/workflows/comfyui-base.yaml b/.github/workflows/comfyui-base.yaml index b1971e32..82281031 100644 --- a/.github/workflows/comfyui-base.yaml +++ b/.github/workflows/comfyui-base.yaml @@ -2,19 +2,13 @@ name: Build and push comfyui-base docker image on: pull_request: - paths: - - docker/Dockerfile.base - - src/comfystream/scripts/ - - configs/ - - .github/workflows/comfyui-base.yaml + paths-ignore: + - "ui/*" branches: - main push: - paths: - - docker/Dockerfile.base - - src/comfystream/scripts/ - - configs/ - - .github/workflows/comfyui-base.yaml + paths-ignore: + - "ui/*" branches: - main tags: diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml new file mode 100644 index 00000000..69e06608 --- /dev/null +++ b/.github/workflows/test.yaml @@ -0,0 +1,86 @@ +name: Test project + +on: + pull_request: + branches: + - "main" + push: + branches: + - "main" + +concurrency: + group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }} + cancel-in-progress: true + +jobs: + codeql: + name: Perform CodeQL analysis + if: ${{ github.repository == 'livepeer/comfystream' }} + runs-on: ubuntu-latest + steps: + - name: Check out code + uses: actions/checkout@v4 + with: + fetch-depth: 0 + ref: ${{ github.event.pull_request.head.sha }} + + # Initializes the CodeQL tools for scanning. + - name: Initialize CodeQL + uses: github/codeql-action/init@v3 + with: + languages: typescript,javascript,python + config-file: ./.github/codeql-config.yaml + + - name: Perform CodeQL Analysis + uses: github/codeql-action/analyze@v3 + + + editorconfig: + name: Run editorconfig checker + runs-on: ubuntu-latest + steps: + - name: Check out code + uses: actions/checkout@v4 + with: + # Check https://github.com/livepeer/go-livepeer/pull/1891 + # for ref value discussion + ref: ${{ github.event.pull_request.head.sha }} + + - name: Install editorconfig-checker + uses: editorconfig-checker/action-editorconfig-checker@main + + - name: Run editorconfig checker against the repo + if: false + # disabled editorconfig lint rule for now + run: editorconfig-checker --format github-actions + + test: + name: Run Tests + runs-on: ubuntu-latest + steps: + - name: Check out code + uses: actions/checkout@v4 + with: + # Check https://github.com/livepeer/go-livepeer/pull/1891 + # for ref value discussion + ref: ${{ github.event.pull_request.head.sha }} + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.11' + cache: pip + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install '.[dev]' + + - name: Run tests + run: pytest --cov --verbose --showlocals + + - name: Upload coverage to Codecov + uses: codecov/codecov-action@v5 + with: + token: ${{ secrets.CI_CODECOV_TOKEN }} + name: ${{ github.event.repository.name }} diff --git a/.gitignore b/.gitignore index 2c8639d3..3c4c381c 100644 --- a/.gitignore +++ b/.gitignore @@ -2,8 +2,8 @@ __pycache__ *.egg-info build .DS_STORE -comfyui/ -ComfyUI/ +./comfyui/ +./ComfyUI/ cuda_temp # VS Code settings @@ -22,3 +22,6 @@ launch.json nodes/web/static/* .cursor/ !nodes/web/static/.gitkeep + +# opencv files +opencv diff --git a/configs/models.yaml b/configs/models.yaml index 115be8ec..4755789d 100644 --- a/configs/models.yaml +++ b/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" \ No newline at end of file + type: "text_encoder" diff --git a/configs/nodes.yaml b/configs/nodes.yaml index 460843a6..38138249 100644 --- a/configs/nodes.yaml +++ b/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 @@ -73,5 +74,5 @@ nodes: comfyui-stream-pack: name: "ComfyUI Stream Pack" url: "https://github.com/livepeer/ComfyUI-Stream-Pack" - branch: "expose_feature_bank_to_comfyui" - type: "utility" \ No newline at end of file + branch: "main" + type: "utility" diff --git a/docker/Dockerfile b/docker/Dockerfile index dd63fd65..10d9ba95 100644 --- a/docker/Dockerfile +++ b/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 diff --git a/docker/Dockerfile.base b/docker/Dockerfile.base index 49a5dd50..00cab374 100644 --- a/docker/Dockerfile.base +++ b/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 diff --git a/docker/entrypoint.sh b/docker/entrypoint.sh index e91dfd97..009f4bda 100755 --- a/docker/entrypoint.sh +++ b/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,18 @@ 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 Dynamic 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/dynamic-dreamshaper8_SD15_\$dyn-b-1-4-2-h-448-704-512-w-448-704-512_00001_.engine \ + --width 512 \ + --height 512 \ + --min-width 448 \ + --min-height 448 \ + --max-width 704 \ + --max-height 704 + + # 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 +71,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 diff --git a/nodes/web/js/comfystream_ui_preview_node.js b/nodes/web/js/comfystream_ui_preview_node.js index 518eb1af..50210313 100644 --- a/nodes/web/js/comfystream_ui_preview_node.js +++ b/nodes/web/js/comfystream_ui_preview_node.js @@ -111,25 +111,10 @@ app.registerExtension({ }); }); - // Update iframe size - this.updateIframeSize(); return result; }; - // Override the resize method to allow both expanding and shrinking - nodeType.prototype.onResize = function(size) { - // Update the size - this.size[0] = size[0]; - this.size[1] = size[1]; - - // Update the iframe size - this.updateIframeSize(); - - // Force canvas update - this.setDirtyCanvas(true, true); - }; - // Add a helper method to update iframe size nodeType.prototype.updateIframeSize = function() { if (this.iframeWidget) { diff --git a/opencv_contrib b/opencv_contrib new file mode 160000 index 00000000..0e5254eb --- /dev/null +++ b/opencv_contrib @@ -0,0 +1 @@ +Subproject commit 0e5254ebf54d2aed6e7eaf6660bf3b797cf50a02 diff --git a/pyproject.toml b/pyproject.toml index 0d9a147d..9240c335 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta" [project] name = "comfystream" description = "Build Live AI Video with ComfyUI" -version = "0.0.4" +version = "0.0.6" license = { file = "LICENSE" } dependencies = [ "asyncio", @@ -15,10 +15,11 @@ dependencies = [ "toml", "twilio", "prometheus_client", + "librosa" ] [project.optional-dependencies] -dev = ["pytest"] +dev = ["pytest", "pytest-cov"] [project.urls] repository = "https://github.com/yondonfu/comfystream" @@ -26,11 +27,11 @@ repository = "https://github.com/yondonfu/comfystream" [tool.comfy] PublisherId = "livepeer-comfystream" DisplayName = "ComfyStream" -Icon = "https://avatars.githubusercontent.com/u/25355022?s=48&v=4" # SVG, PNG, JPG or GIF (MAX. 800x400px) +Icon = "https://raw.githubusercontent.com/livepeer/comfystream-docs/main/logo/icon-light-120px.svg" # SVG, PNG, JPG or GIF (MAX. 800x400px) [tool.setuptools] package-dir = {"" = "src"} packages = {find = {where = ["src", "nodes"]}} [tool.setuptools.dynamic] -dependencies = {file = ["requirements.txt"]} +dependencies = {file = ["requirements.txt"]} \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 4a7e68ad..f94f3345 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,3 +5,4 @@ aiohttp toml twilio prometheus_client +librosa diff --git a/server/app.py b/server/app.py index e1e9150f..83bc943f 100644 --- a/server/app.py +++ b/server/app.py @@ -222,6 +222,9 @@ async def offer(request): pcs.add(pc) tracks = {"video": None, "audio": None} + + # Flag to track if we've received resolution update + resolution_received = {"value": False} # Only add video transceiver if video is present in the offer if "m=video" in offer.sdp: @@ -258,6 +261,27 @@ async def on_message(message): await pipeline.update_prompts(params["prompts"]) response = {"type": "prompts_updated", "success": True} channel.send(json.dumps(response)) + elif params.get("type") == "update_resolution": + if "width" not in params or "height" not in params: + logger.warning("[Control] Missing width or height in update_resolution message") + return + # Update pipeline resolution for future frames + pipeline.width = params["width"] + pipeline.height = params["height"] + logger.info(f"[Control] Updated resolution to {params['width']}x{params['height']}") + + # Mark that we've received resolution + resolution_received["value"] = True + + # Warm the video pipeline with the new resolution + if "m=video" in pc.remoteDescription.sdp: + await pipeline.warm_video() + + response = { + "type": "resolution_updated", + "success": True + } + channel.send(json.dumps(response)) else: logger.warning( "[Server] Invalid message format - missing required fields" @@ -303,10 +327,11 @@ async def on_connectionstatechange(): await pc.setRemoteDescription(offer) + # Only warm audio here, video warming happens after resolution update if "m=audio" in pc.remoteDescription.sdp: await pipeline.warm_audio() - if "m=video" in pc.remoteDescription.sdp: - await pipeline.warm_video() + + # We no longer warm video here - it will be warmed after receiving resolution answer = await pc.createAnswer() await pc.setLocalDescription(answer) @@ -345,7 +370,13 @@ 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' + width=512, + height=512, + 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 +417,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( @@ -435,4 +478,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) diff --git a/server/pipeline.py b/server/pipeline.py index 26270923..d781639e 100644 --- a/server/pipeline.py +++ b/server/pipeline.py @@ -2,24 +2,42 @@ import torch import numpy as np import asyncio +import logging from typing import Any, Dict, Union, List from comfystream.client import ComfyStreamClient +from utils import temporary_log_level WARMUP_RUNS = 5 +logger = logging.getLogger(__name__) + class Pipeline: - def __init__(self, **kwargs): + def __init__(self, width=512, height=512, 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.width = kwargs.get("width", 512) + self.height = kwargs.get("height", 512) + 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): + # Create dummy frame with the CURRENT resolution settings (which might have been updated via control channel) dummy_frame = av.VideoFrame() - dummy_frame.side_data.input = torch.randn(1, 512, 512, 3) + dummy_frame.side_data.input = torch.randn(1, self.height, self.width, 3) + + logger.info(f"Warming video pipeline with resolution {self.width}x{self.height}") for _ in range(WARMUP_RUNS): self.client.put_video_input(dummy_frame) @@ -75,7 +93,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 +109,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 +128,4 @@ async def get_nodes_info(self) -> Dict[str, Any]: return nodes_info async def cleanup(self): - await self.client.cleanup() \ No newline at end of file + await self.client.cleanup() diff --git a/server/utils/__init__.py b/server/utils/__init__.py index 79675ad7..daa71bb1 100644 --- a/server/utils/__init__.py +++ b/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 diff --git a/server/utils/utils.py b/server/utils/utils.py index 63c61a58..c7a7ac30 100644 --- a/server/utils/utils.py +++ b/server/utils/utils.py @@ -6,6 +6,7 @@ import logging from aiohttp import web from typing import List, Tuple +from contextlib import asynccontextmanager logger = logging.getLogger(__name__) @@ -63,3 +64,22 @@ def add_prefix_to_app_routes(app: web.Application, prefix: str): for route in list(app.router.routes()): new_path = prefix + route.resource.canonical app.router.add_route(route.method, new_path, route.handler) + + +@asynccontextmanager +async def temporary_log_level(logger_name: str, level: int): + """Temporarily set the log level of a logger. + + Args: + logger_name: The name of the logger to set the level for. + level: The log level to set. + """ + if level is not None: + logger = logging.getLogger(logger_name) + original_level = logger.level + logger.setLevel(level) + try: + yield + finally: + if level is not None: + logger.setLevel(original_level) diff --git a/src/comfystream/client.py b/src/comfystream/client.py index 6e544eb3..47e995a5 100644 --- a/src/comfystream/client.py +++ b/src/comfystream/client.py @@ -45,7 +45,8 @@ async def run_prompt(self, prompt_index: int): async def cleanup(self): async with self.cleanup_lock: - for task in self.running_prompts.values(): + tasks_to_cancel = list(self.running_prompts.values()) + for task in tasks_to_cancel: task.cancel() try: await task @@ -54,7 +55,11 @@ async def cleanup(self): self.running_prompts.clear() if self.comfy_client.is_running: - await self.comfy_client.__aexit__() + try: + await self.comfy_client.__aexit__() + except Exception as e: + logger.error(f"Error during ComfyClient cleanup: {e}") + await self.cleanup_queues() logger.info("Client cleanup complete") @@ -110,77 +115,105 @@ async def get_available_nodes(self): for node_id, node in prompt.items() } nodes_info = {} - + # Only process nodes until we've found all the ones we need for class_type, node_class in nodes.NODE_CLASS_MAPPINGS.items(): if not remaining_nodes: # Exit early if we've found all needed nodes break - + if class_type not in needed_class_types: continue - + # Get metadata for this node type (same as original get_node_metadata) input_data = node_class.INPUT_TYPES() if hasattr(node_class, 'INPUT_TYPES') else {} input_info = {} - + # Process required inputs if 'required' in input_data: for name, value in input_data['required'].items(): - if isinstance(value, tuple) and len(value) == 2: - input_type, config = value - input_info[name] = { - 'type': input_type, - 'required': True, - 'min': config.get('min', None), - 'max': config.get('max', None), - 'widget': config.get('widget', None) - } + if isinstance(value, tuple): + if len(value) == 1 and isinstance(value[0], list): + # Handle combo box case where value is ([option1, option2, ...],) + input_info[name] = { + 'type': 'combo', + 'value': value[0], # The list of options becomes the value + } + elif len(value) == 2: + input_type, config = value + input_info[name] = { + 'type': input_type, + 'required': True, + 'min': config.get('min', None), + 'max': config.get('max', None), + 'widget': config.get('widget', None) + } + elif len(value) == 1: + # Handle simple type case like ('IMAGE',) + input_info[name] = { + 'type': value[0] + } else: logger.error(f"Unexpected structure for required input {name}: {value}") - - # Process optional inputs + + # Process optional inputs with same logic if 'optional' in input_data: for name, value in input_data['optional'].items(): - if isinstance(value, tuple) and len(value) == 2: - input_type, config = value - input_info[name] = { - 'type': input_type, - 'required': False, - 'min': config.get('min', None), - 'max': config.get('max', None), - 'widget': config.get('widget', None) - } + if isinstance(value, tuple): + if len(value) == 1 and isinstance(value[0], list): + # Handle combo box case where value is ([option1, option2, ...],) + input_info[name] = { + 'type': 'combo', + 'value': value[0], # The list of options becomes the value + } + elif len(value) == 2: + input_type, config = value + input_info[name] = { + 'type': input_type, + 'required': False, + 'min': config.get('min', None), + 'max': config.get('max', None), + 'widget': config.get('widget', None) + } + elif len(value) == 1: + # Handle simple type case like ('IMAGE',) + input_info[name] = { + 'type': value[0] + } else: logger.error(f"Unexpected structure for optional input {name}: {value}") - + # Now process any nodes in our prompt that use this class_type for node_id in list(remaining_nodes): node = prompt[node_id] if node.get('class_type') != class_type: continue - + node_info = { 'class_type': class_type, 'inputs': {} } - + if 'inputs' in node: for input_name, input_value in node['inputs'].items(): + input_metadata = input_info.get(input_name, {}) node_info['inputs'][input_name] = { 'value': input_value, - 'type': input_info.get(input_name, {}).get('type', 'unknown'), - 'min': input_info.get(input_name, {}).get('min', None), - 'max': input_info.get(input_name, {}).get('max', None), - 'widget': input_info.get(input_name, {}).get('widget', None) + 'type': input_metadata.get('type', 'unknown'), + 'min': input_metadata.get('min', None), + 'max': input_metadata.get('max', None), + 'widget': input_metadata.get('widget', None) } - + # For combo type inputs, include the list of options + if input_metadata.get('type') == 'combo': + node_info['inputs'][input_name]['value'] = input_metadata.get('value', []) + nodes_info[node_id] = node_info remaining_nodes.remove(node_id) all_prompts_nodes_info[prompt_index] = nodes_info - + return all_prompts_nodes_info - + except Exception as e: logger.error(f"Error getting node info: {str(e)}") return {} diff --git a/src/comfystream/scripts/build_trt.py b/src/comfystream/scripts/build_trt.py index 69caf199..fa0a0f24 100644 --- a/src/comfystream/scripts/build_trt.py +++ b/src/comfystream/scripts/build_trt.py @@ -50,7 +50,7 @@ def parse_args(): parser = argparse.ArgumentParser( - description="Build a static TensorRT engine from a ComfyUI model." + description="Build a TensorRT engine from a ComfyUI model." ) parser.add_argument( "--model", @@ -74,14 +74,21 @@ def parse_args(): "--width", type=int, default=512, - help="Width in pixels for the exported model (default 1024)", + help="Width in pixels for the exported model (default 512)", ) parser.add_argument( "--height", type=int, default=512, - help="Height in pixels for the exported model (default 1024)", + help="Height in pixels for the exported model (default 512)", ) + + # Dynamic Engine Optional Args + parser.add_argument("--min-width", type=int, default=None, help="Minimum width for dynamic shape (optional)") + parser.add_argument("--min-height", type=int, default=None, help="Minimum height for dynamic shape (optional)") + parser.add_argument("--max-width", type=int, default=None, help="Maximum width for dynamic shape (optional)") + parser.add_argument("--max-height", type=int, default=None, help="Maximum height for dynamic shape (optional)") + parser.add_argument( "--context", type=int, @@ -101,12 +108,16 @@ def parse_args(): ) return parser.parse_args() -def build_static_trt_engine( +def build_trt_engine( model_path: str, engine_out_path: str, batch_size_opt: int = 1, width_opt: int = 512, height_opt: int = 512, + min_width: int = None, + min_height: int = None, + max_width: int = None, + max_height: int = None, context_opt: int = 1, num_video_frames: int = 14, fp8: bool = False, @@ -114,8 +125,8 @@ def build_static_trt_engine( ): """ 1) Load the model from ComfyUI by path or name - 2) Export to ONNX (static shape) - 3) Build a static TensorRT .engine file + 2) Export to ONNX + 3) Build a TensorRT .engine file """ # Check if the engine file already exists @@ -171,22 +182,32 @@ def build_static_trt_engine( fp8 = fp8, ) - # 3) Build the static TRT engine + # 3) Build the TRT engine model_version = detect_version_from_model(loaded_model) model_helper = get_helper_from_model(loaded_model) trt_model = TRTDiffusionBackbone(model_helper) - # We'll define min/opt/max config all the same (i.e. 'static') - # TODO: make this configurable + # Dynamic engine support: only if min/max width/height provided + is_dynamic = all(v is not None for v in [min_width, max_width, min_height, max_height]) min_config = { + "batch_size": batch_size_opt, + "height": min_height if is_dynamic else height_opt, + "width": min_width if is_dynamic else width_opt, + "context_len": context_opt * model_helper.context_len, + } + opt_config = { "batch_size": batch_size_opt, "height": height_opt, "width": width_opt, "context_len": context_opt * model_helper.context_len, } - opt_config = dict(min_config) - max_config = dict(min_config) + max_config = { + "batch_size": batch_size_opt, + "height": max_height if is_dynamic else height_opt, + "width": max_width if is_dynamic else width_opt, + "context_len": context_opt * model_helper.context_len, + } # The tensorrt_diffusion_model build() signature is typically: # build(onnx_path, engine_path, timing_cache_path, opt_config, min_config, max_config) @@ -223,12 +244,16 @@ def build_static_trt_engine( def main(): args = parse_args() - build_static_trt_engine( + build_trt_engine( model_path = args.model, engine_out_path = args.out_engine, batch_size_opt = args.batch_size, height_opt = args.height, width_opt = args.width, + min_width = args.min_width, + min_height = args.min_height, + max_width = args.max_width, + max_height = args.max_height, context_opt = args.context, fp8 = args.fp8, verbose = args.verbose diff --git a/ui/package-lock.json b/ui/package-lock.json index d26fdca3..9abcb69f 100644 --- a/ui/package-lock.json +++ b/ui/package-lock.json @@ -1,12 +1,12 @@ { "name": "ui", - "version": "0.0.4", + "version": "0.0.6", "lockfileVersion": 3, "requires": true, "packages": { "": { "name": "ui", - "version": "0.0.4", + "version": "0.0.6", "dependencies": { "@hookform/resolvers": "^3.9.1", "@radix-ui/react-dialog": "^1.1.6", @@ -17,7 +17,7 @@ "class-variance-authority": "^0.7.0", "clsx": "^2.1.1", "lucide-react": "^0.454.0", - "next": "15.1.6", + "next": "15.2.4", "next-themes": "^0.4.4", "react": "^19.0.0", "react-dom": "^19.0.0", @@ -713,9 +713,9 @@ } }, "node_modules/@next/env": { - "version": "15.1.6", - "resolved": "https://registry.npmjs.org/@next/env/-/env-15.1.6.tgz", - "integrity": "sha512-d9AFQVPEYNr+aqokIiPLNK/MTyt3DWa/dpKveiAaVccUadFbhFEvY6FXYX2LJO2Hv7PHnLBu2oWwB4uBuHjr/w==", + "version": "15.2.4", + "resolved": "https://registry.npmjs.org/@next/env/-/env-15.2.4.tgz", + "integrity": "sha512-+SFtMgoiYP3WoSswuNmxJOCwi06TdWE733D+WPjpXIe4LXGULwEaofiiAy6kbS0+XjM5xF5n3lKuBwN2SnqD9g==", "license": "MIT" }, "node_modules/@next/eslint-plugin-next": { @@ -729,9 +729,9 @@ } }, "node_modules/@next/swc-darwin-arm64": { - "version": "15.1.6", - "resolved": "https://registry.npmjs.org/@next/swc-darwin-arm64/-/swc-darwin-arm64-15.1.6.tgz", - "integrity": "sha512-u7lg4Mpl9qWpKgy6NzEkz/w0/keEHtOybmIl0ykgItBxEM5mYotS5PmqTpo+Rhg8FiOiWgwr8USxmKQkqLBCrw==", + "version": "15.2.4", + "resolved": "https://registry.npmjs.org/@next/swc-darwin-arm64/-/swc-darwin-arm64-15.2.4.tgz", + "integrity": "sha512-1AnMfs655ipJEDC/FHkSr0r3lXBgpqKo4K1kiwfUf3iE68rDFXZ1TtHdMvf7D0hMItgDZ7Vuq3JgNMbt/+3bYw==", "cpu": [ "arm64" ], @@ -745,9 +745,9 @@ } }, "node_modules/@next/swc-darwin-x64": { - "version": "15.1.6", - "resolved": "https://registry.npmjs.org/@next/swc-darwin-x64/-/swc-darwin-x64-15.1.6.tgz", - "integrity": "sha512-x1jGpbHbZoZ69nRuogGL2MYPLqohlhnT9OCU6E6QFewwup+z+M6r8oU47BTeJcWsF2sdBahp5cKiAcDbwwK/lg==", + "version": "15.2.4", + "resolved": "https://registry.npmjs.org/@next/swc-darwin-x64/-/swc-darwin-x64-15.2.4.tgz", + "integrity": "sha512-3qK2zb5EwCwxnO2HeO+TRqCubeI/NgCe+kL5dTJlPldV/uwCnUgC7VbEzgmxbfrkbjehL4H9BPztWOEtsoMwew==", "cpu": [ "x64" ], @@ -761,9 +761,9 @@ } }, "node_modules/@next/swc-linux-arm64-gnu": { - "version": "15.1.6", - "resolved": "https://registry.npmjs.org/@next/swc-linux-arm64-gnu/-/swc-linux-arm64-gnu-15.1.6.tgz", - "integrity": "sha512-jar9sFw0XewXsBzPf9runGzoivajeWJUc/JkfbLTC4it9EhU8v7tCRLH7l5Y1ReTMN6zKJO0kKAGqDk8YSO2bg==", + "version": "15.2.4", + "resolved": "https://registry.npmjs.org/@next/swc-linux-arm64-gnu/-/swc-linux-arm64-gnu-15.2.4.tgz", + "integrity": "sha512-HFN6GKUcrTWvem8AZN7tT95zPb0GUGv9v0d0iyuTb303vbXkkbHDp/DxufB04jNVD+IN9yHy7y/6Mqq0h0YVaQ==", "cpu": [ "arm64" ], @@ -777,9 +777,9 @@ } }, "node_modules/@next/swc-linux-arm64-musl": { - "version": "15.1.6", - "resolved": "https://registry.npmjs.org/@next/swc-linux-arm64-musl/-/swc-linux-arm64-musl-15.1.6.tgz", - "integrity": "sha512-+n3u//bfsrIaZch4cgOJ3tXCTbSxz0s6brJtU3SzLOvkJlPQMJ+eHVRi6qM2kKKKLuMY+tcau8XD9CJ1OjeSQQ==", + "version": "15.2.4", + "resolved": "https://registry.npmjs.org/@next/swc-linux-arm64-musl/-/swc-linux-arm64-musl-15.2.4.tgz", + "integrity": "sha512-Oioa0SORWLwi35/kVB8aCk5Uq+5/ZIumMK1kJV+jSdazFm2NzPDztsefzdmzzpx5oGCJ6FkUC7vkaUseNTStNA==", "cpu": [ "arm64" ], @@ -793,9 +793,9 @@ } }, "node_modules/@next/swc-linux-x64-gnu": { - "version": "15.1.6", - "resolved": "https://registry.npmjs.org/@next/swc-linux-x64-gnu/-/swc-linux-x64-gnu-15.1.6.tgz", - "integrity": "sha512-SpuDEXixM3PycniL4iVCLyUyvcl6Lt0mtv3am08sucskpG0tYkW1KlRhTgj4LI5ehyxriVVcfdoxuuP8csi3kQ==", + "version": "15.2.4", + "resolved": "https://registry.npmjs.org/@next/swc-linux-x64-gnu/-/swc-linux-x64-gnu-15.2.4.tgz", + "integrity": "sha512-yb5WTRaHdkgOqFOZiu6rHV1fAEK0flVpaIN2HB6kxHVSy/dIajWbThS7qON3W9/SNOH2JWkVCyulgGYekMePuw==", "cpu": [ "x64" ], @@ -809,9 +809,9 @@ } }, "node_modules/@next/swc-linux-x64-musl": { - "version": "15.1.6", - "resolved": "https://registry.npmjs.org/@next/swc-linux-x64-musl/-/swc-linux-x64-musl-15.1.6.tgz", - "integrity": "sha512-L4druWmdFSZIIRhF+G60API5sFB7suTbDRhYWSjiw0RbE+15igQvE2g2+S973pMGvwN3guw7cJUjA/TmbPWTHQ==", + "version": "15.2.4", + "resolved": "https://registry.npmjs.org/@next/swc-linux-x64-musl/-/swc-linux-x64-musl-15.2.4.tgz", + "integrity": "sha512-Dcdv/ix6srhkM25fgXiyOieFUkz+fOYkHlydWCtB0xMST6X9XYI3yPDKBZt1xuhOytONsIFJFB08xXYsxUwJLw==", "cpu": [ "x64" ], @@ -825,9 +825,9 @@ } }, "node_modules/@next/swc-win32-arm64-msvc": { - "version": "15.1.6", - "resolved": "https://registry.npmjs.org/@next/swc-win32-arm64-msvc/-/swc-win32-arm64-msvc-15.1.6.tgz", - "integrity": "sha512-s8w6EeqNmi6gdvM19tqKKWbCyOBvXFbndkGHl+c9YrzsLARRdCHsD9S1fMj8gsXm9v8vhC8s3N8rjuC/XrtkEg==", + "version": "15.2.4", + "resolved": "https://registry.npmjs.org/@next/swc-win32-arm64-msvc/-/swc-win32-arm64-msvc-15.2.4.tgz", + "integrity": "sha512-dW0i7eukvDxtIhCYkMrZNQfNicPDExt2jPb9AZPpL7cfyUo7QSNl1DjsHjmmKp6qNAqUESyT8YFl/Aw91cNJJg==", "cpu": [ "arm64" ], @@ -841,9 +841,9 @@ } }, "node_modules/@next/swc-win32-x64-msvc": { - "version": "15.1.6", - "resolved": "https://registry.npmjs.org/@next/swc-win32-x64-msvc/-/swc-win32-x64-msvc-15.1.6.tgz", - "integrity": "sha512-6xomMuu54FAFxttYr5PJbEfu96godcxBTRk1OhAvJq0/EnmFU/Ybiax30Snis4vdWZ9LGpf7Roy5fSs7v/5ROQ==", + "version": "15.2.4", + "resolved": "https://registry.npmjs.org/@next/swc-win32-x64-msvc/-/swc-win32-x64-msvc-15.2.4.tgz", + "integrity": "sha512-SbnWkJmkS7Xl3kre8SdMF6F/XDh1DTFEhp0jRTj/uB8iPKoU2bb2NDfcu+iifv1+mxQEd1g2vvSxcZbXSKyWiQ==", "cpu": [ "x64" ], @@ -5013,12 +5013,12 @@ "license": "MIT" }, "node_modules/next": { - "version": "15.1.6", - "resolved": "https://registry.npmjs.org/next/-/next-15.1.6.tgz", - "integrity": "sha512-Hch4wzbaX0vKQtalpXvUiw5sYivBy4cm5rzUKrBnUB/y436LGrvOUqYvlSeNVCWFO/770gDlltR9gqZH62ct4Q==", + "version": "15.2.4", + "resolved": "https://registry.npmjs.org/next/-/next-15.2.4.tgz", + "integrity": "sha512-VwL+LAaPSxEkd3lU2xWbgEOtrM8oedmyhBqaVNmgKB+GvZlCy9rgaEc+y2on0wv+l0oSFqLtYD6dcC1eAedUaQ==", "license": "MIT", "dependencies": { - "@next/env": "15.1.6", + "@next/env": "15.2.4", "@swc/counter": "0.1.3", "@swc/helpers": "0.5.15", "busboy": "1.6.0", @@ -5033,14 +5033,14 @@ "node": "^18.18.0 || ^19.8.0 || >= 20.0.0" }, "optionalDependencies": { - "@next/swc-darwin-arm64": "15.1.6", - "@next/swc-darwin-x64": "15.1.6", - "@next/swc-linux-arm64-gnu": "15.1.6", - "@next/swc-linux-arm64-musl": "15.1.6", - "@next/swc-linux-x64-gnu": "15.1.6", - "@next/swc-linux-x64-musl": "15.1.6", - "@next/swc-win32-arm64-msvc": "15.1.6", - "@next/swc-win32-x64-msvc": "15.1.6", + "@next/swc-darwin-arm64": "15.2.4", + "@next/swc-darwin-x64": "15.2.4", + "@next/swc-linux-arm64-gnu": "15.2.4", + "@next/swc-linux-arm64-musl": "15.2.4", + "@next/swc-linux-x64-gnu": "15.2.4", + "@next/swc-linux-x64-musl": "15.2.4", + "@next/swc-win32-arm64-msvc": "15.2.4", + "@next/swc-win32-x64-msvc": "15.2.4", "sharp": "^0.33.5" }, "peerDependencies": { diff --git a/ui/package.json b/ui/package.json index 1e6a4e14..e3b01ba6 100644 --- a/ui/package.json +++ b/ui/package.json @@ -1,6 +1,6 @@ { "name": "ui", - "version": "0.0.4", + "version": "0.0.6", "private": true, "scripts": { "dev": "cross-env NEXT_PUBLIC_DEV=true next dev", @@ -21,7 +21,7 @@ "class-variance-authority": "^0.7.0", "clsx": "^2.1.1", "lucide-react": "^0.454.0", - "next": "15.1.6", + "next": "15.2.4", "next-themes": "^0.4.4", "react": "^19.0.0", "react-dom": "^19.0.0", diff --git a/ui/src/components/control-panel.tsx b/ui/src/components/control-panel.tsx index 8d9a90f6..3d5f458c 100644 --- a/ui/src/components/control-panel.tsx +++ b/ui/src/components/control-panel.tsx @@ -4,7 +4,7 @@ import React, { useState, useEffect } from "react"; import { usePeerContext } from "@/context/peer-context"; import { usePrompt } from "./settings"; -type InputValue = string | number | boolean; +type InputValue = string | number | boolean | string[]; interface InputInfo { value: InputValue; @@ -12,6 +12,7 @@ interface InputInfo { min?: number; max?: number; widget?: string; + options?: string[]; } interface NodeInfo { @@ -45,19 +46,32 @@ const InputControl = ({ value: string; onChange: (value: string) => void; }) => { - if (input.widget === "combo") { + if (input.widget === "combo" || input.type === "combo") { + // Get options from either the options field or value field + const options = input.options + ? input.options + : Array.isArray(input.value) + ? input.value + : typeof input.value === 'string' + ? [input.value] + : []; + + // If no value is selected, select the first option by default + const currentValue = value || options[0] || ''; + return ( ); } @@ -100,6 +114,9 @@ const InputControl = ({ className="p-2 border rounded w-full" /> ); + // Handle combo in the main combo block above + case "combo": + return InputControl({ input: { ...input, widget: "combo" }, value, onChange }); default: console.warn(`Unhandled input type: ${input.type}`); // Debug log return ( @@ -196,12 +213,33 @@ export const ControlPanel = ({ ] : null; if (!currentInput || !currentPrompts) return; + + // Don't send updates if this is a combo and we haven't selected a value yet + if (currentInput.widget === "combo" && !panelState.value) return; let isValidValue = true; let processedValue: InputValue = panelState.value; - // Validate and process value based on type - switch (currentInput.type.toLowerCase()) { + // For combo inputs, use the value directly + if (currentInput.widget === "combo" || currentInput.type === "combo") { + // Get options from either the options field or value field + const options = currentInput.options + ? currentInput.options + : Array.isArray(currentInput.value) + ? currentInput.value as string[] + : typeof currentInput.value === 'string' + ? [currentInput.value as string] + : []; + + // If no value is selected and we have options, use the first option + const validValue = panelState.value || options[0] || ''; + + // Validate that the value is in the options list + isValidValue = options.includes(validValue); + processedValue = validValue; + } else { + // Validate and process value based on type + switch (currentInput.type.toLowerCase()) { case "number": isValidValue = /^-?\d*\.?\d*$/.test(panelState.value) && panelState.value !== ""; @@ -217,13 +255,9 @@ export const ControlPanel = ({ processedValue = panelState.value; break; default: - if (currentInput.widget === "combo") { - isValidValue = panelState.value !== ""; - processedValue = panelState.value; - } else { - isValidValue = panelState.value !== ""; - processedValue = panelState.value; - } + isValidValue = panelState.value !== ""; + processedValue = panelState.value; + } } const hasRequiredFields = @@ -261,9 +295,32 @@ export const ControlPanel = ({ } const updatedPrompt = JSON.parse(JSON.stringify(prompt)); // Deep clone if (updatedPrompt[panelState.nodeId]?.inputs) { - updatedPrompt[panelState.nodeId].inputs[panelState.fieldName] = - processedValue; - hasUpdated = true; + // Ensure we're not overwriting with an invalid value + const currentVal = updatedPrompt[panelState.nodeId].inputs[panelState.fieldName]; + const input = availableNodes[promptIdxToUpdate][panelState.nodeId]?.inputs[panelState.fieldName]; + + if (input?.widget === 'combo' || input?.type === 'combo') { + // Get options from either the options field or value field + const options = input.options + ? input.options + : Array.isArray(input.value) + ? input.value as string[] + : typeof input.value === 'string' + ? [input.value as string] + : []; + + // If no value is selected and we have options, use the first option + const validValue = (processedValue as string) || options[0] || ''; + + // Only update if it's a valid combo value + if (options.includes(validValue)) { + updatedPrompt[panelState.nodeId].inputs[panelState.fieldName] = validValue; + hasUpdated = true; + } + } else { + updatedPrompt[panelState.nodeId].inputs[panelState.fieldName] = processedValue; + hasUpdated = true; + } } return updatedPrompt; }, @@ -312,8 +369,13 @@ export const ControlPanel = ({ if (input.type.toLowerCase() === "boolean") { return (!!input.value).toString(); } - if (input.widget === "combo" && Array.isArray(input.value)) { - return input.value[0]?.toString() || ""; + if (input.widget === "combo") { + const options = Array.isArray(input.value) + ? input.value as string[] + : typeof input.value === 'string' + ? [input.value as string] + : []; + return options[0] || ""; } return input.value?.toString() || "0"; }; @@ -327,11 +389,19 @@ export const ControlPanel = ({ selectedField ]; if (input) { - const initialValue = getInitialValue(input); - onStateChange({ - fieldName: selectedField, - value: initialValue, - }); + // For combo fields, don't set an initial value to prevent auto-update from firing + if (input.widget === "combo") { + onStateChange({ + fieldName: selectedField, + value: "", + }); + } else { + const initialValue = getInitialValue(input); + onStateChange({ + fieldName: selectedField, + value: initialValue, + }); + } } else { onStateChange({ fieldName: selectedField }); } @@ -357,7 +427,7 @@ export const ControlPanel = ({ onStateChange({ nodeId: e.target.value, fieldName: "", - value: "0", + value: "", // Start with empty value to prevent auto-update from firing }); }} className="p-2 border rounded" @@ -387,16 +457,20 @@ export const ControlPanel = ({ typeof info.type === "string" ? info.type.toLowerCase() : String(info.type).toLowerCase(); - return ( - ["boolean", "number", "float", "int", "string"].includes( - type, - ) || info.widget === "combo" - ); + return [ + "boolean", + "number", + "float", + "int", + "string", + "combo", + ].includes( + type, + ) || info.widget === "combo"; }) .map(([field, info]) => ( ))} diff --git a/ui/src/components/peer.tsx b/ui/src/components/peer.tsx index dae60781..bfcc2d47 100644 --- a/ui/src/components/peer.tsx +++ b/ui/src/components/peer.tsx @@ -10,6 +10,10 @@ export interface PeerProps extends React.HTMLAttributes { onConnected: () => void; onDisconnected: () => void; localStream: MediaStream | null; + resolution?: { + width: number; + height: number; + }; } export const PeerConnector = (props: PeerProps) => { diff --git a/ui/src/components/room.tsx b/ui/src/components/room.tsx index de6d94d0..0c4b3876 100644 --- a/ui/src/components/room.tsx +++ b/ui/src/components/room.tsx @@ -1,7 +1,7 @@ "use client"; import { PeerConnector } from "@/components/peer"; -import { StreamConfig, StreamSettings } from "@/components/settings"; +import { StreamConfig, StreamSettings, DEFAULT_CONFIG } from "@/components/settings"; import { Webcam } from "@/components/webcam"; import { usePeerContext } from "@/context/peer-context"; import { useCallback, useEffect, useRef, useState } from "react"; @@ -14,13 +14,50 @@ import { } from "@/components/ui/tooltip"; import { ControlPanelsContainer } from "@/components/control-panels-container"; +// Custom hook for managing toast lifecycle +function useToast() { + const toastIdRef = useRef(undefined); + + const showToast = useCallback((message: string, type: 'loading' | 'success' | 'error' = 'loading') => { + // Always dismiss previous toast first + if (toastIdRef.current) { + toast.dismiss(toastIdRef.current); + } + + // Create new toast based on type + let id; + if (type === 'loading') { + id = toast.loading(message); + } else if (type === 'success') { + id = toast.success(message); + } else if (type === 'error') { + id = toast.error(message); + } + + toastIdRef.current = id; + return id; + }, []); + + const dismissToast = useCallback(() => { + if (toastIdRef.current) { + toast.dismiss(toastIdRef.current); + toastIdRef.current = undefined; + } + }, []); + + return { showToast, dismissToast, toastId: toastIdRef }; +} + interface MediaStreamPlayerProps { stream: MediaStream; + resolution: { width: number; height: number }; } -function MediaStreamPlayer({ stream }: MediaStreamPlayerProps) { +function MediaStreamPlayer({ stream, resolution, onFrame }: MediaStreamPlayerProps & { onFrame?: () => void }) { const videoRef = useRef(null); const [needsPlayButton, setNeedsPlayButton] = useState(false); + const frameCheckRef = useRef(null); + const lastTimeRef = useRef(0); useEffect(() => { if (!videoRef.current || !stream) return; @@ -29,6 +66,29 @@ function MediaStreamPlayer({ stream }: MediaStreamPlayerProps) { video.srcObject = stream; setNeedsPlayButton(false); + // Add frame detection if needed + if (onFrame) { + // Use requestAnimationFrame for more frequent checks + const checkFrame = (time: number) => { + // If the video's currentTime has changed, we have a new frame + if (video.currentTime !== lastTimeRef.current) { + lastTimeRef.current = video.currentTime; + onFrame(); + } + + frameCheckRef.current = requestAnimationFrame(checkFrame); + }; + + frameCheckRef.current = requestAnimationFrame(checkFrame); + + return () => { + if (frameCheckRef.current !== null) { + cancelAnimationFrame(frameCheckRef.current); + frameCheckRef.current = null; + } + }; + } + // Handle autoplay const playStream = async () => { try { @@ -50,8 +110,13 @@ function MediaStreamPlayer({ stream }: MediaStreamPlayerProps) { video.srcObject = null; video.pause(); } + + if (frameCheckRef.current !== null) { + cancelAnimationFrame(frameCheckRef.current); + frameCheckRef.current = null; + } }; - }, [stream]); + }, [stream, onFrame]); const handlePlayClick = async () => { try { @@ -65,12 +130,15 @@ function MediaStreamPlayer({ stream }: MediaStreamPlayerProps) { }; return ( -
+