fix: merge conflicts

.gitignore

@@ -2,4 +2,17 @@ __pycache__
 *.egg-info
 build
 .DS_STORE
-comfyui*
+comfyui*
+
+# VS Code settings
+.vscode/
+*.code-workspace
+launch.json
+
+# Environment files
+.env
+.env.local
+.env.*.local
+.env.development
+.env.test
+.env.production

README.md

@@ -5,16 +5,23 @@ comfystream is a package for running img2img [Comfy](https://www.comfy.org/) wor
 This repo also includes a WebRTC server and UI that uses comfystream to support streaming from a webcam and processing the stream with a workflow JSON file (API format) created in ComfyUI. If you have an existing ComfyUI installation, the same custom nodes used to create the workflow in ComfyUI will be re-used when processing the video stream.
 
 - [comfystream](#comfystream)
-- [Install package](#install-package)
-  - [Custom Nodes](#custom-nodes)
-  - [Usage](#usage)
-- [Run tests](#run-tests)
-- [Run server](#run-server)
-- [Run UI](#run-ui)
-- [Limitations](#limitations)
-- [Troubleshoot](#troubleshoot)
+  - [Quick Start](#quick-start)
+  - [Install package](#install-package)
+    - [Custom Nodes](#custom-nodes)
+    - [Usage](#usage)
+  - [Run tests](#run-tests)
+  - [Run server](#run-server)
+  - [Run UI](#run-ui)
+  - [Limitations](#limitations)
+  - [Troubleshoot](#troubleshoot)
 
-# Install package 
+## Quick Start
+
+The fastest way to get started is to follow [this tutorial](https://livepeer.notion.site/ComfyStream-Dev-Environment-Setup-15d0a3485687802e9528d26050142d82) by @ryanontheinside.
+
+For additional information, refer to the remaining sections below.
+
+## Install package
 
 **Prerequisites**
 
@@ -24,21 +31,21 @@ A separate environment can be used to avoid any dependency issues with an existi
 
 Create the environment:
 
-```
+```bash
 conda create -n comfystream python=3.11
 ```
 
 Activate the environment:
 
-```
+```bash
 conda activate comfystream
 ```
 
 Make sure you have [PyTorch](https://pytorch.org/get-started/locally/) installed.
 
 Install `comfystream`:
 
-```
+```bash
 pip install git+https://github.com/yondonfu/comfystream.git
 
 # This can be used to install from a local repo
@@ -47,75 +54,75 @@ pip install git+https://github.com/yondonfu/comfystream.git
 # pip install -e .
 ```
 
-## Custom Nodes
+### Custom Nodes
 
 comfystream uses a few custom nodes to support running workflows.
 
 Copy the custom nodes into the `custom_nodes` folder of your ComfyUI workspace:
 
-```
+```bash
 cp -r nodes/* custom_nodes/
 ```
 
 For example, if your ComfyUI workspace is under `/home/user/ComfyUI`:
 
-```
+```bash
 cp -r nodes/* /home/user/ComfyUI/custom_nodes
 ```
 
-## Usage
+### Usage
 
 See `example.py`.
 
-# Run tests
+## Run tests
 
 Install dev dependencies:
 
-```
+```bash
 pip install .[dev]
 ```
 
 Run tests:
 
-```
+```bash
 pytest
 ```
 
-# Run server
+## Run server
 
 Install dependencies:
 
-```
+```bash
 pip install -r requirements.txt
 ```
 
 If you have existing custom nodes in your ComfyUI workspace, you will need to install their requirements in your current environment:
 
-```
+```bash
 python install.py --workspace <COMFY_WORKSPACE>
 ```
 
 Run the server:
 
-```
+```bash
 python server/app.py --workspace <COMFY_WORKSPACE>
 ```
 
 Show additional options for configuring the server:
 
-```
+```bash
 python server/app.py -h
 ```
 
 **Remote Setup**
 
 A local server should connect with a local UI out-of-the-box. It is also possible to run a local UI and connect with a remote server, but there may be additional dependencies.
 
-In order for the remote server to connect with another peer (i.e. a browser) without any additional dependencies you will need to allow inbound/outbound UDP traffic on ports 1024-65535 ([source](https://github.com/aiortc/aiortc/issues/490#issuecomment-788807118)). 
+In order for the remote server to connect with another peer (i.e. a browser) without any additional dependencies you will need to allow inbound/outbound UDP traffic on ports 1024-65535 ([source](https://github.com/aiortc/aiortc/issues/490#issuecomment-788807118)).
 
 If you only have a subset of those UDP ports available, you can use the `--media-ports` flag to specify a comma delimited list of ports to use:
 
-```
+```bash
 python server/app.py --workspace <COMFY_WORKSPACE> --media-ports 1024,1025,...
 ```
 
@@ -124,38 +131,38 @@ If you are running the server in a restrictive network environment where this is
 At the moment, the server supports using Twilio's TURN servers (although it is easy to make the update to support arbitrary TURN servers):
 
 1. Sign up for a [Twilio](https://www.twilio.com/en-us) account.
-2. Copy the Account SID and Auth Token from https://console.twilio.com/.
+2. Copy the Account SID and Auth Token from [https://console.twilio.com/](https://console.twilio.com/).
 3. Set the `TWILIO_ACCOUNT_SID` and `TWILIO_AUTH_TOKEN` environment variables.
 
-````
+```bash
 export TWILIO_ACCOUNT_SID=...
 export TWILIO_AUTH_TOKEN=...
-````
+```
 
-# Run UI
+## Run UI
 
 **Prerequisities**
 
 - [Node.js](https://nodejs.org/en/download/package-manager)
 
 Install dependencies
 
-```
+```bash
 cd ui
 npm install --legacy-peer-deps
 ```
 
 Run local dev server:
 
-```
+```bash
 npm run dev
 ```
 
-By default the app will be available at http://localhost:3000.
+By default the app will be available at <http://localhost:3000>.
 
-The Stream URL is the URL of the [server](#run-server) which defaults to http://127.0.0.1:8888.
+The Stream URL is the URL of the [server](#run-server) which defaults to <http://127.0.0.1:8888>.
 
-# Limitations
+## Limitations
 
 At the moment, a workflow must fufill the following requirements:
 
@@ -170,12 +177,12 @@ At the moment, a workflow must fufill the following requirements:
 - The workflow must have a single output using a PreviewImage or SaveImage node
   - At runtime, this node is replaced with a SaveTensor node
 
-# Troubleshoot
+## Troubleshoot
 
 This project has been tested locally successfully with the following setup:
 
 - OS: Ubuntu
 - GPU: Nvidia RTX 4090
 - Driver: 550.127.05
 - CUDA: 12.5
-- torch: 2.5.1+cu121
+- torch: 2.5.1+cu121

install.py

@@ -3,9 +3,9 @@
 import argparse
 import logging
 
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 logger = logging.getLogger(__name__)
 
-
 def install_custom_node_req(workspace: str):
     custom_nodes_path = os.path.join(workspace, "custom_nodes")
     for folder in os.listdir(custom_nodes_path):
@@ -24,4 +24,6 @@ def install_custom_node_req(workspace: str):
     )
     args = parser.parse_args()
 
+    logger.info("Installing custom node requirements...")
     install_custom_node_req(args.workspace)
+    logger.info("Custom node requirements installed successfully.")

server/app.py

@@ -12,6 +12,7 @@
     RTCConfiguration,
     RTCIceServer,
     MediaStreamTrack,
+    RTCDataChannel,
 )
 from aiortc.rtcrtpsender import RTCRtpSender
 from aiortc.codecs import h264
@@ -20,6 +21,7 @@
 
 logger = logging.getLogger(__name__)
 
+
 MAX_BITRATE = 2000000
 MIN_BITRATE = 2000000
 
@@ -132,6 +134,39 @@ async def offer(request):
     h264.MAX_BITRATE = MAX_BITRATE
     h264.MIN_BITRATE = MIN_BITRATE
 
+    # Handle control channel from client
+    @pc.on("datachannel")
+    def on_datachannel(channel):
+        if channel.label == "control":
+            @channel.on("message")
+            async def on_message(message):
+                try:
+                    params = json.loads(message)
+
+                    if params.get("type") == "get_nodes":
+                        nodes_info = await pipeline.get_nodes_info()
+                        response = {
+                            "type": "nodes_info",
+                            "nodes": nodes_info
+                        }
+                        channel.send(json.dumps(response))
+                    elif params.get("type") == "update_prompt":
+                        if "prompt" not in params:
+                            logger.warning("[Control] Missing prompt in update_prompt message")
+                            return
+                        pipeline.set_prompt(params["prompt"])
+                        response = {
+                            "type": "prompt_updated",
+                            "success": True
+                        }
+                        channel.send(json.dumps(response))
+                    else:
+                        logger.warning("[Server] Invalid message format - missing required fields")
+                except json.JSONDecodeError:
+                    logger.error("[Server] Invalid JSON received")
+                except Exception as e:
+                    logger.error(f"[Server] Error processing message: {str(e)}")
+
     @pc.on("track")
     def on_track(track):
         logger.info(f"Track received: {track.kind}")
@@ -222,7 +257,11 @@ async def on_shutdown(app: web.Application):
     )
     args = parser.parse_args()
 
-    logging.basicConfig(level=args.log_level.upper())
+    logging.basicConfig(
+        level=args.log_level.upper(),
+        format='%(asctime)s [%(levelname)s] %(message)s',
+        datefmt='%H:%M:%S'
+    )
 
     app = web.Application()
     app["media_ports"] = args.media_ports.split(",") if args.media_ports else None

server/pipeline.py

@@ -11,6 +11,8 @@
 
 WARMUP_RUNS = 5
 
+
+
 class Pipeline:
     def __init__(self, **kwargs):
         self.client = ComfyStreamClient(**kwargs, max_workers=5) # hardcoded max workers
@@ -26,6 +28,9 @@ def __init__(self, **kwargs):
         self.time_base = fractions.Fraction(1, self.sample_rate)
         self.curr_pts = 0 # figure out a better way to set back pts to processed audio frames
 
+    def set_prompt(self, prompt: Dict[Any, Any]):
+        self.client.set_prompt(prompt)
+
     async def warm(self):
         dummy_video_frame = torch.randn(1, 512, 512, 3)
         dummy_audio_frame = np.random.randint(-32768, 32767, 48000 * 1, dtype=np.int16)
@@ -95,7 +100,6 @@ async def get_processed_video_frame(self):
         frame.time_base = time_base
         return frame
 
-
     async def get_processed_audio_frame(self):
         while not self.audio_output_frames:
             out_fut = await self.audio_futures.get()
@@ -104,4 +108,9 @@ async def get_processed_audio_frame(self):
                 print("No Audio output")
                 continue
             self.audio_output_frames.extend(self.audio_postprocess(output))
-        return self.audio_output_frames.pop(0)
+        return self.audio_output_frames.pop(0)
+
+    async def get_nodes_info(self) -> Dict[str, Any]:
+        """Get information about all nodes in the current prompt including metadata."""
+        nodes_info = await self.client.get_available_nodes()
+        return nodes_info