feat: comfystream v0.0.4 ai-runner port (#479)

- Improved pipeline flexibility and async handling.
- Fixed environment, event loop, and CI issues.
- Updated base images and cleaned dependencies.
- Added frame processing and warm video interface.
- Reverted and optimized configurations for stability.

---------

Co-authored-by: Elite Encoder <john@eliteencoder.net>
Co-authored-by: Peter Schroedl <peter_schroedl@me.com>
Co-authored-by: Rafal Leszko <rafal@livepeer.org>

Author: 4 people

.devcontainer/devcontainer.json

@@ -1,34 +1,52 @@
 // For format details, see https://aka.ms/devcontainer.json. For config options, see the
 // README at: https://github.com/devcontainers/templates/tree/main/src/ubuntu
 {
-	"name": "ai-runner",
-	// Image to use for the dev container. More info: https://containers.dev/guide/dockerfile.
-	"build": {
-		"dockerfile": "../Dockerfile",
-		// "dockerfile": "../docker/Dockerfile.text_to_speech",
-		"context": ".."
-	},
-	"runArgs": [
-		"--gpus=all"
-	],
-	// Features to add to the dev container. More info: https://containers.dev/features.
-	// Configure tool-specific properties.
-	"customizations": {
-		"vscode": {
-			"settings": {},
-			"extensions": [
-				"ms-python.python",
-				"ms-python.black-formatter"
-			]
-		}
-	},
-	// Use 'forwardPorts' to make a list of ports inside the container available locally.
-	"forwardPorts": [
-		8000
-	],
-	// Use 'mounts' to make a list of local folders available inside the container.
-	"mounts": [
-		// "source=${localWorkspaceFolder}/models,target=/models,type=bind"
-		"source=${localEnv:HOME}/.lpData/models,target=/models,type=bind"
-	]
+        "name": "ai-runner",
+        "initializeCommand": "ls",
+        // Image to use for the dev container. More info: https://containers.dev/guide/dockerfile.
+        "containerEnv": {
+            "PIPELINE": "comfyui"
+        },
+        "build": {
+            "dockerfile": "../runner/docker/Dockerfile.live-app__PIPELINE__",
+            "args": {
+                "PIPELINE": "comfyui"
+            },
+            // "dockerfile": "../Dockerfile",
+            // "dockerfile": "../docker/Dockerfile.text_to_speech",
+            "context": "../runner"
+        },
+        "runArgs": [
+            "--gpus=all"
+        ],
+        // Features to add to the dev container. More info: https://containers.dev/features.
+        // Configure tool-specific properties.
+        "customizations": {
+            "vscode": {
+                "settings": {
+                    "python.defaultInterpreterPath": "/workspace/miniconda3/envs/comfystream/bin/python",
+                    "python.venvPath": "/workspace/miniconda3/envs",
+                    "python.terminal.activateEnvInCurrentTerminal": false,
+                    "python.terminal.activateEnvironment": true,
+                    "terminal.integrated.shellIntegration.enabled": true
+                },
+                "extensions": [
+                    "ms-python.python",
+                    "ms-python.black-formatter"
+                ]
+            }
+        },
+        // Use 'forwardPorts' to make a list of ports inside the container available locally.
+        "forwardPorts": [
+            8000
+        ],
+        "appPort": [
+            "8000:8000"
+        ],
+        // Use 'mounts' to make a list of local folders available inside the container.
+        "mounts": [
+            // "source=${localWorkspaceFolder}/models,target=/models,type=bind"
+            "source=${localEnv:HOME}/models/ComfyUI--models/,target=/workspace/ComfyUI/models,type=bind",
+            "source=${localEnv:HOME}/models/ComfyUI--output/,target=/workspace/ComfyUI/output,type=bind"
+        ]
 }

runner/app/live/pipelines/comfyui.py

@@ -1,18 +1,20 @@
 import os
 import json
 import torch
-from PIL import Image
 import asyncio
 import numpy as np
+from PIL import Image
 from typing import Union
 from pydantic import BaseModel, field_validator
 
 from .interface import Pipeline
 from comfystream.client import ComfyStreamClient
+from trickle import VideoFrame, VideoOutput
 
 import logging
 
 COMFY_UI_WORKSPACE_ENV = "COMFY_UI_WORKSPACE"
+WARMUP_RUNS = 1
 DEFAULT_WORKFLOW_JSON = json.loads("""
 {
   "1": {
@@ -263,46 +265,54 @@ def validate_prompt(cls, v) -> dict:
 
 
 class ComfyUI(Pipeline):
-    def __init__(self, **params):
-        super().__init__(**params)
-
+    def __init__(self):
         comfy_ui_workspace = os.getenv(COMFY_UI_WORKSPACE_ENV)
         self.client = ComfyStreamClient(cwd=comfy_ui_workspace)
         self.params: ComfyUIParams
+        self.video_incoming_frames = asyncio.Queue()
 
-        self.update_params(**params)
+    async def initialize(self, **params):
+        new_params = ComfyUIParams(**params)
+        logging.info(f"Initializing ComfyUI Pipeline with prompt: {new_params.prompt}")
+        # TODO: currently its a single prompt, but need to support multiple prompts
+        await self.client.set_prompts([new_params.prompt])
+        self.params = new_params
+
+        # Warm up the pipeline
+        dummy_frame = VideoFrame(None, 0, 0)
+        dummy_frame.side_data.input = torch.randn(1, 512, 512, 3)
 
-        # Comfy will cache nodes that only need to be run once (i.e. a node that loads model weights)
-        # We can run the prompt once before actual inputs come in to "warmup"
-        warmup_input = torch.randn(1, 512, 512, 3)
-        asyncio.get_event_loop().run_until_complete(self.client.queue_prompt(warmup_input))
+        for _ in range(WARMUP_RUNS):
+            self.client.put_video_input(dummy_frame)
+            _ = await self.client.get_video_output()
+        logging.info("Pipeline initialization and warmup complete")
 
-    def process_frame(self, image: Image.Image) -> Image.Image:
-        # Normalize by dividing by 255 to ensure the tensor values are between 0 and 1
-        image_np = np.array(image.convert("RGB")).astype(np.float32) / 255.0
-        # Convert from numpy to torch.Tensor
-        # Initially, the torch.Tensor will have shape HWC but we want BHWC
-        # unsqueeze(0) will add a batch dimension at the beginning of 1 which means we just have 1 image
-        image_tensor = torch.tensor(image_np).unsqueeze(0)
+    async def put_video_frame(self, frame: VideoFrame):
+        image_np = np.array(frame.image.convert("RGB")).astype(np.float32) / 255.0
+        frame.side_data.input = torch.tensor(image_np).unsqueeze(0)
+        frame.side_data.skipped = True
+        self.client.put_video_input(frame)
+        await self.video_incoming_frames.put(frame)
 
-        # Process using ComfyUI pipeline
-        result_tensor = asyncio.get_event_loop().run_until_complete(self.client.queue_prompt(image_tensor))
+    async def get_processed_video_frame(self, request_id):
+        result_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()
 
-        # Convert back from Tensor to PIL.Image
         result_tensor = result_tensor.squeeze(0)
         result_image_np = (result_tensor * 255).byte()
         result_image = Image.fromarray(result_image_np.cpu().numpy())
-        return result_image
-
-    def update_params(self, **params):
+        return VideoOutput(frame.replace_image(result_image), request_id)
+    
+    async def update_params(self, **params):
         new_params = ComfyUIParams(**params)
-        logging.info(f"ComfyUI Pipeline Prompt: {new_params.prompt}")
-        self.client.set_prompt(new_params.prompt)
+        logging.info(f"Updating ComfyUI Pipeline Prompt: {new_params.prompt}")
+        # TODO: currently its a single prompt, but need to support multiple prompts
+        await self.client.update_prompts([new_params.prompt])
         self.params = new_params
 
-    #TODO: This is a hack to stop the ComfyStreamClient. Use the comfystream api to stop the client in 0.0.2
     async def stop(self):
         logging.info("Stopping ComfyUI pipeline")
-        if self.client.comfy_client.is_running:
-            await self.client.comfy_client.__aexit__(None, None, None)
+        await self.client.stop()
         logging.info("ComfyUI pipeline stopped")

runner/app/live/pipelines/interface.py

@@ -1,5 +1,6 @@
-from abc import ABC, abstractmethod
 from PIL import Image
+from abc import ABC, abstractmethod
+from trickle import VideoFrame, VideoOutput
 
 class Pipeline(ABC):
     """Abstract base class for image processing pipelines.
@@ -22,21 +23,38 @@ def __init__(self, **params):
         pass
 
     @abstractmethod
-    def process_frame(self, frame: Image.Image) -> Image.Image:
-        """Process a single frame through the pipeline.
-
-        Called sequentially with each frame from the stream.
+    async def put_video_frame(self, frame: VideoFrame):
+        """Put a frame into the pipeline.
 
         Args:
-            frame: Input PIL Image
+            frame: Input VideoFrame
+        """
+        pass
+
+    @abstractmethod
+    async def get_processed_video_frame(self, request_id: str = '') -> VideoOutput:
+        """Get a processed frame from the pipeline.
 
         Returns:
-            Processed PIL Image
+            Processed VideoFrame
+        """
+        pass
+
+    @abstractmethod
+    async def initialize(self, **params):
+        """Initialize the pipeline with parameters and warm up the processing.
+
+        This method sets up the initial pipeline state and performs warmup operations. 
+        Must maintain valid state on success or restore previous state on failure.
+        Starts the pipeline loops in comfystream.
+
+        Args:
+            **params: Implementation-specific parameters
         """
         pass
 
     @abstractmethod
-    def update_params(self, **params):
+    async def update_params(self, **params):
         """Update pipeline parameters.
 
         Must maintain valid state on success or restore previous state on failure.

runner/app/live/pipelines/loader.py

@@ -1,10 +1,10 @@
 from .interface import Pipeline
 
-def load_pipeline(name: str, **params) -> Pipeline:
+def load_pipeline(name: str) -> Pipeline:
     if name == "comfyui":
         from .comfyui import ComfyUI
-        return ComfyUI(**params)
+        return ComfyUI()
     elif name == "noop":
         from .noop import Noop
-        return Noop(**params)
+        return Noop()
     raise ValueError(f"Unknown pipeline: {name}")

runner/app/live/pipelines/noop.py

@@ -1,14 +1,29 @@
-from PIL import Image
 import logging
+import asyncio
+from PIL import Image
+
 
 from .interface import Pipeline
+from trickle import VideoFrame, VideoOutput
 
 class Noop(Pipeline):
-  def __init__(self, **params):
-    super().__init__(**params)
+  def __init__(self):
+    self.frame_queue = asyncio.Queue()
 
-  def process_frame(self, image: Image.Image) -> Image.Image:
-    return image.convert("RGB")
+  async def put_video_frame(self, frame: VideoFrame):
+    await self.frame_queue.put(frame)
 
-  def update_params(self, **params):
+  async def get_processed_video_frame(self) -> VideoOutput:
+    frame = await self.frame_queue.get()
+    processed_frame = frame.image.convert("RGB")
+    return VideoOutput(frame.replace_image(processed_frame))
+  
+  async def initialize(self, **params):
+    logging.info(f"Initializing Noop pipeline with params: {params}")
+    logging.info("Pipeline initialization complete")
+
+  async def update_params(self, **params):
     logging.info(f"Updating params: {params}")
+
+  async def stop(self):
+    logging.info("Stopping pipeline")