Merge branch 'main' into main

Author: eliteprox

.gitignore

@@ -2,8 +2,9 @@ __pycache__
 *.egg-info
 build
 .DS_STORE
-comfyui*
-ComfyUI*
+comfyui/
+ComfyUI/
+cuda_temp
 
 # VS Code settings
 .vscode/
@@ -18,4 +19,4 @@ launch.json
 .env.test
 .env.production
 
-nodes/web/static/*
+nodes/web/static/*

README.md

@@ -10,6 +10,7 @@ This repo also includes a WebRTC server and UI that uses comfystream to support
     - [Docker Image](#docker-image)
       - [RunPod](#runpod)
       - [Tensordock](#tensordock)
+      - [Other Cloud Providers](#other-cloud-providers)
   - [Download Models](#download-models)
   - [Install package](#install-package)
     - [Custom Nodes](#custom-nodes)
@@ -42,11 +43,15 @@ The RunPod template [livepeer-comfystream](https://runpod.io/console/deploy?temp
 
 #### Tensordock
 
-We also have a python script that can be used to spin up a ComfyStream instance on a [Tensordock server](https://tensordock.com/). Refer to [scripts/README.md](scripts/README.md) for instructions.
+We also have a python script that can be used to spin up a ComfyStream instance on a [Tensordock server](https://tensordock.com/). Refer to [scripts/README.md](./scripts/README.md#tensordock-auto-setup-fully-automated) for instructions.
+
+#### Other Cloud Providers
+
+We also provide an [Ansible playbook](https://docs.ansible.com/ansible/latest/installation_guide/index.html) for deploying ComfyStream on any cloud provider. Refer to [scripts/README.md](./scripts/README.md#cloud-agnostic-automated-setup-ansible-based-deployment) for instructions.
 
 ## Download Models
 
-Refer to [src/comfystream/scripts/README.md](src/comfystream/scripts/README.md) for instructions to download commonly used models.
+Refer to [scripts/README.md](src/comfystream/scripts/README.md) for instructions to download commonly used models.
 
 ## Install package
 
@@ -83,7 +88,7 @@ pip install git+https://github.com/yondonfu/comfystream.git
 
 ### Custom Nodes
 
-comfystream uses a few auxiliary custom nodes to support running workflows. 
+Comfystream uses a few auxiliary custom nodes to support running workflows.
 
 **Note:** If you are using comfystream as a custom node in ComfyUI, you can skip the following steps.
 

configs/nodes.yaml

@@ -2,9 +2,12 @@ nodes:
   # Core TensorRT nodes
   comfyui-tensorrt:
     name: "ComfyUI TensorRT"
-    url: "https://github.com/yondonfu/ComfyUI_TensorRT"
-    branch: "quantization_with_controlnet_fixes"
+    url: "https://github.com/eliteprox/ComfyUI_TensorRT.git"
+    branch: "fix/onnxconverter-converter"
     type: "tensorrt"
+    dependencies: 
+      - "onnxruntime-gpu>=1.17.0"
+      - "onnx>=1.17.0"
 
   comfyui-depthanything-tensorrt:
     name: "ComfyUI DepthAnything TensorRT"
@@ -65,6 +68,4 @@ nodes:
   comfyui_controlnet_aux:
     name: "ComfyUI ControlNet Auxiliary"
     url: "https://github.com/Fannovel16/comfyui_controlnet_aux"
-    type: "utility"
-    dependencies: 
-      - "mediapipe==0.10.8"
+    type: "utility"

docker/Dockerfile.base

@@ -52,12 +52,11 @@ COPY	./workflows/comfyui/*	/workspace/ComfyUI/user/default/workflows
 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 && \
-	conda run -n comfystream --no-capture-output --cwd /workspace/comfystream pip install . --root-user-action=ignore && \
-	ln -s /workspace/comfystream /workspace/ComfyUI/custom_nodes/comfystream && \
-	conda run -n comfystream --no-capture-output --cwd /workspace/comfystream python install.py --workspace /workspace/ComfyUI && \
-	conda run -n comfystream --no-capture-output pip install --upgrade tensorrt-cu12-bindings tensorrt-cu12-libs --root-user-action=ignore && \
-	conda run -n comfystream --no-capture-output pip install mediapipe==0.10.8 --root-user-action=ignore
+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 python install.py --workspace /workspace/ComfyUI
+RUN conda run -n comfystream --no-capture-output pip install --upgrade tensorrt-cu12-bindings tensorrt-cu12-libs --root-user-action=ignore
 
 # Configure no environment activation by default
 RUN	conda config --set auto_activate_base false && \

install.py

@@ -53,6 +53,7 @@ def download_and_extract_ui_files(version: str):
     output_dir = os.path.join(os.getcwd(), "nodes", "web")
     pathlib.Path(output_dir).mkdir(parents=True, exist_ok=True)
     base_url = urllib.parse.urljoin("https://github.com/yondonfu/comfystream/releases/download/", f"v{version}/static.tar.gz")
+    fallback_url = "https://github.com/yondonfu/comfystream/releases/latest/download/static.tar.gz"
 
     # Create a temporary directory instead of a temporary file
     with tempfile.TemporaryDirectory() as temp_dir:
@@ -63,13 +64,25 @@ def download_and_extract_ui_files(version: str):
         logger.info(f"Downloading {base_url}")
         try:
             urllib.request.urlretrieve(base_url, download_path)
-            
-            # Extract contents
+        except urllib.error.HTTPError as e:
+            if e.code == 404:
+                logger.warning(f"{base_url} not found, trying {fallback_url}")
+                try:
+                    urllib.request.urlretrieve(fallback_url, download_path)
+                except Exception as e:
+                    logger.error(f"Error downloading latest ui package: {e}")
+                    raise
+            else:
+                logger.error(f"Error downloading package: {e}")
+                raise
+        
+        # Extract contents
+        try:
             logger.info(f"Extracting files to {output_dir}")
             with tarfile.open(download_path, 'r:gz') as tar:
                 tar.extractall(path=output_dir)
         except Exception as e:
-            logger.error(f"Error downloading or extracting files: {e}")
+            logger.error(f"Error extracting files: {e}")
             raise
 
 if __name__ == "__main__":
@@ -100,14 +113,14 @@ def download_and_extract_ui_files(version: str):
 
     if workspace is None:
         logger.warning("No ComfyUI workspace found. Please specify a valid workspace path to fully install")
-
-    logger.info("Downloading and extracting UI files...")
-    version = get_project_version(os.getcwd())
-    download_and_extract_ui_files(version)
 
     if workspace is not None:
         logger.info("Ensuring __init__.py files exist in ComfyUI directories...")
         ensure_init_files(workspace)
         logger.info("Installing custom node requirements...")
     subprocess.check_call([sys.executable, "-m", "pip", "install", "-e", "."])
+    
+    logger.info("Downloading and extracting UI files...")
+    version = get_project_version(os.getcwd())
+    download_and_extract_ui_files(version)
     logger.info("Installation completed successfully.")

nodes/tensor_utils/prestartup_script.py

@@ -28,26 +28,29 @@ def wrapped_control_merge(self, control, control_prev, output_dtype):
                 }
 
             # Get result from original merge function
-            result = original_control_merge(self, control, control_prev, output_dtype)
-           
-            # Clone all output tensors
-            result = {
-                k: [t.clone() if t is not None else None for t in v]
-                for k, v in result.items()
-            }
-           
-            # Mark CUDA graph step at end
-            if torch.cuda.is_available() and hasattr(torch.compiler, 'cudagraph_mark_step_begin'):
-                torch.compiler.cudagraph_mark_step_begin()
-                torch.cuda.synchronize()
-               
-            return result
-           
+            try:
+                result = original_control_merge(self, control, control_prev, output_dtype)
+            
+                # Clone all output tensors
+                result = {
+                    k: [t.clone() if t is not None else None for t in v]
+                    for k, v in result.items()
+                }
+                return result
+            except Exception as e:
+                print(f"Error: Failed to clone ControlNet during inference: {str(e)}")
+                raise e
+            finally:
+                # Mark CUDA graph step at end
+                if torch.cuda.is_available() and hasattr(torch.compiler, 'cudagraph_mark_step_begin'):
+                    torch.compiler.cudagraph_mark_step_begin()
+                    torch.cuda.synchronize()
+
         # Apply the patch
         ControlBase.control_merge = wrapped_control_merge
         print("Successfully patched ControlNet for torch.compile() compatibility")
     except Exception as e:
         print(f"Warning: Failed to patch ControlNet: {str(e)}")
 
 # Apply patch when module is imported
-patch_controlnet_for_stream() 
+patch_controlnet_for_stream() 

scripts/README.md

@@ -1,48 +1,112 @@
 # ComfyUI Helper Scripts
 
-This directory contains various helper scripts designed to streamline working with ComfyStream:
+This directory contains helper scripts to simplify the deployment and management of **ComfyStream**:
 
-- `spinup_comfystream_tensordock.py`: Spins up a ComfyStream instance on a [Tensordock server](https://tensordock.com/), installs dependencies, and starts the server.
+- **Ansible Playbook (`ansible/plays/setup_comfystream.yml`)** – Deploys ComfyStream on any cloud provider.  
+- **`spinup_comfystream_tensordock.py`** – Fully automates VM creation and ComfyStream setup on a [TensorDock server](https://tensordock.com/).  
 
 ## Usage Instructions
 
-### Spinup TensorDock ComfyStream Instance
+### Ansible Playbook (Cloud-Agnostic Deployment)
 
-1. **Create a Tensordock Account**: Register at [Tensordock](https://dashboard.tensordock.com/register), attach a credit card, and create an API key and token.
-2. **Set Up a Python Virtual Environment**:
-   Create and activate a virtual environment using [Conda](https://docs.anaconda.com/miniconda/) and install the required dependencies:
+This repository includes an [Ansible playbook](https://docs.ansible.com/ansible/latest/playbook_guide/playbooks_intro.html) to deploy ComfyStream on any cloud provider. Follow the steps below:
+
+1. **Create a VM**:  
+   Deploy a VM on a cloud provider such as [TensorDock](https://marketplace.tensordock.com/deploy?gpu=geforcertx4090-pcie-24gb&gpuCount=1&ramAmount=16&vcpuCount=4&storage=100&os=Ubuntu-22.04-LTS), AWS, Google Cloud, or Azure with the following minimum specifications:
+   - **GPU**: 20GB VRAM  
+   - **RAM**: 16GB  
+   - **CPU**: 4 vCPUs  
+   - **Storage**: 100GB
+
+   > [!TIP]
+   > You can use the `spinup_comfystream_tensordock.py --bare-vm` script to create a [compatible VM on TensorDock](https://marketplace.tensordock.com/deploy?gpu=geforcertx4090-pcie-24gb&gpuCount=1&ramAmount=16&vcpuCount=4&storage=100&os=Ubuntu-22.04-LTS)
+
+2. **Open Required Ports**:  
+   Ensure the following ports are open **inbound and outbound** on the VM's firewall/security group:
+   - **SSH (Port 22)** – Remote access  
+   - **HTTPS (Port 8189)** – ComfyStream access  
+  
+3. **Install Ansible**:  
+   Follow the [official Ansible installation guide](https://docs.ansible.com/ansible/latest/installation_guide/index.html).
+
+4. **Configure the Inventory File**:
+   Add the VM’s public IP to `ansible/inventory.yml`.
+
+5. **Change the ComfyUI Password:**  
+   Open the `ansible/plays/setup_comfystream.yaml` file and replace the `comfyui_password` value with your own secure password.
+
+6. **Run the Playbook**:  
+   Execute:
 
     ```bash
-    conda create -n comfystream python=3.8
-    conda activate comfystream
-    pip install -r requirements.txt
+    ansible-playbook -i ansible/inventory.yaml ansible/plays/setup_comfystream.yaml
     ```
 
-3. **Check Available Options** *(Optional but Recommended)*:  
-   To see all available script options, run:
+   > [!IMPORTANT]  
+   > When using a non-sudo user, add `--ask-become-pass` to provide the sudo password or use an Ansible vault for secure storage.
+
+7. **Access the Server**:  
+   After the playbook completes, **ComfyStream** will start, and you can access **ComfyUI** at `https://<VM_IP>:<PORT_FOR_8189>`. Credentials are shown in the output and regenerated each time. To persist the password, set the `comfyui_password` variable when running the playbook:
+
+   ```bash
+   ansible-playbook -i ansible/inventory.yaml ansible/plays/setup_comfystream.yaml -e "comfyui_password=YourSecurePasswordHere"
+   ```
+
+> [!IMPORTANT]
+> If you encounter a `toomanyrequests` error while pulling the Docker image, either wait a few minutes or provide your Docker credentials when running the playbook:  
+>
+> ```bash
+> ansible-playbook -i ansible/inventory.yaml ansible/plays/setup_comfystream.yaml -e "docker_hub_username=your_dockerhub_username docker_hub_password=your_dockerhub_pat"
+> ```
+
+> [!TIP]
+> The [ComfyStream Docker image](https://hub.docker.com/r/livepeer/comfystream/tags) is **~20GB**. To check download progress, SSH into the VM and run:
+>
+> ```bash
+> docker pull livepeer/comfystream:latest
+> ```
+
+### TensorDock Spinup Script (Fully Automated)
+
+The `spinup_comfystream_tensordock.py` script automates VM provisioning, setup, and server launch on [TensorDock](https://tensordock.com/). Follow the steps below:
+
+1. **Create a TensorDock Account**: Sign up at [Tensordock](https://dashboard.tensordock.com/register), add a payment method, and generate API credentials.
+
+2. **Set Up a Python Virtual Environment**:
+   To prevent dependency conflicts, create and activate a virtual environment with [Conda](https://docs.anaconda.com/miniconda/) and install the required dependencies:
+
+    ```bash
+   conda create -n comfystream python=3.8
+   conda activate comfystream
+   pip install -r requirements.txt
+   ```
+
+3. **View Available Script Options** *(Optional)*:  
+   To see all available options, run:
 
     ```bash
     python spinup_comfystream_tensordock.py --help
     ```
 
 4. **Run the Script**:  
-   Execute the following command to spin up a ComfyStream instance:
+   Execute the following command to provision a VM and set up ComfyStream automatically:
 
     ```bash
     python spinup_comfystream_tensordock.py --api-key <API_KEY> --api-token <API_TOKEN>
     ```
 
 5. **Access the Server**:  
-   The script will set up the instance, install dependencies, and start the server. This process will take a few minutes. Once completed, you can access the ComfyStream server and UI using the provided URLs.
+   Once the setup is complete, the script will display the URLs to access ComfyStream.
 
-6. **Stop the Server**:  
-   To stop the server, run:
+6. **Stop & Delete the VM** *(When No Longer Needed)*:
+   To stop and remove the instance, run:
 
     ```bash
     python spinup_comfystream_tensordock.py --delete <VM_ID>
     ```
 
-   Replace `<VM_ID>` with the ID of the VM you want to delete. You can find the VM ID in the script logs or the [Tensordock dashboard](https://dashboard.tensordock.com/).
+   Replace `<VM_ID>` with the VM ID found in the script logs or the [TensorDock dashboard](https://dashboard.tensordock.com/instances).
 
 > [!WARNING]
-> If you see `max retries exceeded with url` errors, the VM was likely created but is inaccessible. Check the [TensorDock dashboard](https://dashboard.tensordock.com/instances), delete the VM, wait 2-3 minutes, then run the script again.
+> If you encounter `max retries exceeded with URL` errors, the VM might have been created but is inaccessible.  
+> Check the [TensorDock dashboard](https://dashboard.tensordock.com/instances), delete the VM manually, wait **2-3 minutes**, then rerun the script.

scripts/ansible/inventory.yaml

@@ -0,0 +1,8 @@
+---
+all:
+  hosts:
+    comfystream_vm:
+      ansible_host: 123.234.102.170                     # Replace with your VM's IP address
+      ansible_port: 22                                  # Replace with your VM's SSH port
+      ansible_user: user                                # Replace with your SSH user
+      ansible_ssh_private_key_file: ~/.ssh/private_key  # Path to your SSH key