ci: Add mypy for typing

.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"
+        ]
 }

.github/workflows/ai-runner-live-pipelines-docker.yaml

@@ -39,7 +39,7 @@ jobs:
         uses: tj-actions/changed-files@823fcebdb31bb35fdf2229d9f769b400309430d0 # v46.0.3
         with:
           files: |
-            runner/docker/Dockerfile.live-base
+            runner/docker/Dockerfile.live-base*
 
       - name: Check if build needed
         id: check_build

.github/workflows/ai-runner-test.yaml

@@ -30,6 +30,10 @@ jobs:
           pip install -r runner/requirements-dev.txt
           pip install -r runner/requirements.txt
 
+      - name: Run mypy type checking
+        working-directory: runner
+        run: mypy app/
+
       - name: Run tests
         working-directory: runner
         run: pytest --verbose --showlocals

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 warm_video(self):
+        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("Video frame warmup done")
 
-    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):
+        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
+        return frame.replace_image(result_image)
+
+    async def set_params(self, **params):
+        new_params = ComfyUIParams(**params)
+        logging.info(f"Setting ComfyUI Pipeline 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
 
-    def update_params(self, **params):
+    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,37 @@ 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) -> VideoFrame:
+        """Get a processed frame from the pipeline.
 
         Returns:
-            Processed PIL Image
+            Processed VideoFrame
+        """
+        pass
+
+    @abstractmethod
+    async def set_params(self, **params):
+        """Set pipeline parameters initally.
+
+        Must maintain valid state on success or restore previous state on failure.
+        set_params starts the prompt 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,31 @@
-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()
+
+  async def put_video_frame(self, frame: VideoFrame):
+    await self.frame_queue.put(frame)
 
-  def process_frame(self, image: Image.Image) -> Image.Image:
-    return image.convert("RGB")
+  async def get_processed_video_frame(self) -> VideoFrame:
+    frame = await self.frame_queue.get()
+    processed_frame = frame.image.convert("RGB")
+    return frame.replace_image(processed_frame)
 
-  def update_params(self, **params):
+  async def warm_video(self):
+    logging.info("Warming video")
+
+  async def set_params(self, **params):
+    logging.info(f"Setting params: {params}")
+
+  async def update_params(self, **params):
     logging.info(f"Updating params: {params}")
+
+  async def stop(self):
+    logging.info("Stopping pipeline")