feat(pipeline): ensure warmup uses generated frames instead of input frames (#544)

* feat(pipeline): enhance warmup handling and default workflow application

- Added `ensure_warmup` method to the Pipeline class to manage warmup state based on resolution.
- Integrated warmup checks in the ComfyStreamFrameProcessor to ensure the pipeline is ready before processing streams.
- Implemented default workflow application when no prompts are provided during stream initialization.
- Improved error handling for warmup failures and streamlined frame processing logic.

* feat(frame_processor): add loading overlay management during stream processing to avoid blocking frame processor

- Implemented methods to toggle and manage a loading overlay in the ComfyStreamFrameProcessor.
- Added functionality to disable the overlay after pipeline ingest resumes.
- Enhanced error handling for loading overlay state updates.
- Integrated loading overlay management into the stream start process to improve user experience.

* fix(frame_processor) fix imports after merge

* bump pytrickle

* bump pytrickle dependency to version 0.1.6 in pyproject.toml

* bump pytrickle dependency to version 0.1.7 in pyproject.toml

Author: eliteprox

pyproject.toml

@@ -22,7 +22,7 @@ dependencies = [
 [project.optional-dependencies]
 dev = ["pytest", "pytest-cov", "ruff"]
 server = [
-    "pytrickle @ git+https://github.com/livepeer/[email protected].5"
+    "pytrickle @ git+https://github.com/livepeer/[email protected].7"
 ]
 
 [project.urls]

server/frame_processor.py

@@ -11,7 +11,10 @@
 
 from comfystream.pipeline import Pipeline
 from comfystream.pipeline_state import PipelineState
-from comfystream.utils import convert_prompt
+from comfystream.utils import (
+    convert_prompt,
+    get_default_workflow,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -151,6 +154,49 @@ async def _forward_text_loop():
         except Exception:
             logger.warning("Failed to set up text monitoring", exc_info=True)
 
+    def _set_loading_overlay(self, enabled: bool) -> bool:
+        """Toggle the StreamProcessor loading overlay if available."""
+        processor = self._stream_processor
+        if not processor:
+            return False
+        try:
+            processor.set_loading_overlay(enabled)
+            logger.debug("Set loading overlay to %s", enabled)
+            return True
+        except Exception:
+            logger.warning("Failed to update loading overlay state", exc_info=True)
+            return False
+
+    def _schedule_overlay_reset_on_ingest_enabled(self) -> None:
+        """Disable the loading overlay after pipeline ingest resumes."""
+        if not self.pipeline:
+            self._set_loading_overlay(False)
+            return
+
+        if self.pipeline.is_ingest_enabled():
+            self._set_loading_overlay(False)
+            return
+
+        async def _wait_for_ingest_enable():
+            try:
+                while True:
+                    if self._stop_event.is_set():
+                        break
+                    if not self.pipeline:
+                        break
+                    if self.pipeline.is_ingest_enabled():
+                        break
+                    await asyncio.sleep(0.05)
+            except asyncio.CancelledError:
+                raise
+            except Exception:
+                logger.debug("Loading overlay watcher error", exc_info=True)
+            finally:
+                self._set_loading_overlay(False)
+
+        task = asyncio.create_task(_wait_for_ingest_enable())
+        self._background_tasks.append(task)
+
     async def _stop_text_forwarder(self) -> None:
         """Stop the background text forwarder task if running."""
         task = self._text_forward_task
@@ -212,16 +258,27 @@ async def on_stream_start(self, params: Optional[Dict[str, Any]] = None):
         logger.info("Stream starting")
         self._reset_stop_event()
         logger.info(f"Stream start params: {params}")
+        overlay_managed = False
 
         if not self.pipeline:
             logger.debug("Stream start requested before pipeline initialization")
             return
 
         stream_params = normalize_stream_params(params)
+        stream_width = stream_params.get("width")
+        stream_height = stream_params.get("height")
+        stream_width = int(stream_width) if stream_width is not None else None
+        stream_height = int(stream_height) if stream_height is not None else None
         prompt_payload = stream_params.pop("prompts", None)
         if prompt_payload is None:
             prompt_payload = stream_params.pop("prompt", None)
 
+        if not prompt_payload and not self.pipeline.state_manager.is_initialized():
+            logger.info(
+                "No prompts provided for new stream; applying default workflow for initialization"
+            )
+            prompt_payload = get_default_workflow()
+
         if prompt_payload:
             try:
                 await self._apply_stream_start_prompt(prompt_payload)
@@ -240,6 +297,19 @@ async def on_stream_start(self, params: Optional[Dict[str, Any]] = None):
                 logger.exception("Failed to process stream start parameters")
                 return
 
+        overlay_managed = self._set_loading_overlay(True)
+
+        try:
+            await self.pipeline.ensure_warmup(stream_width, stream_height)
+        except Exception:
+            if overlay_managed:
+                self._set_loading_overlay(False)
+            logger.exception("Failed to ensure pipeline warmup during stream start")
+            return
+
+        if overlay_managed:
+            self._schedule_overlay_reset_on_ingest_enabled()
+
         try:
             if (
                 self.pipeline.state != PipelineState.STREAMING
@@ -312,6 +382,12 @@ async def process_video_async(
             if not self.pipeline:
                 return frame
 
+            # TODO: Do we really need this here?
+            await self.pipeline.ensure_warmup()
+
+            if not self.pipeline.state_manager.is_initialized():
+                return VideoProcessingResult.WITHHELD
+
             # If pipeline ingestion is paused, withhold frame so pytrickle renders the overlay
             if not self.pipeline.is_ingest_enabled():
                 return VideoProcessingResult.WITHHELD
@@ -324,18 +400,9 @@ async def process_video_async(
             # Process through pipeline
             await self.pipeline.put_video_frame(av_frame)
 
-            # Try to get processed frame with short timeout
-            try:
-                processed_av_frame = await asyncio.wait_for(
-                    self.pipeline.get_processed_video_frame(),
-                    timeout=self._stream_processor.overlay_config.auto_timeout_seconds,
-                )
-                processed_frame = VideoFrame.from_av_frame_with_timing(processed_av_frame, frame)
-                return processed_frame
-
-            except asyncio.TimeoutError:
-                # No frame ready yet - return withheld sentinel to trigger overlay
-                return VideoProcessingResult.WITHHELD
+            processed_av_frame = await self.pipeline.get_processed_video_frame()
+            processed_frame = VideoFrame.from_av_frame_with_timing(processed_av_frame, frame)
+            return processed_frame
 
         except Exception as e:
             logger.error(f"Video processing failed: {e}")

src/comfystream/pipeline.py

@@ -72,6 +72,10 @@ def __init__(
         self._initialize_lock = asyncio.Lock()
         self._ingest_enabled = True
         self._prompt_update_lock = asyncio.Lock()
+        self._warmup_lock = asyncio.Lock()
+        self._warmup_task: Optional[asyncio.Task] = None
+        self._warmup_completed = False
+        self._last_warmup_resolution: Optional[tuple[int, int]] = None
 
     @property
     def state(self) -> PipelineState:
@@ -155,6 +159,10 @@ async def warmup(
             await self.state_manager.transition_to(PipelineState.ERROR)
             raise
         finally:
+            if warmup_successful:
+                self._warmup_completed = True
+                self._last_warmup_resolution = (self.width, self.height)
+
             if transitioned and warmup_successful:
                 try:
                     await self.state_manager.transition_to(PipelineState.READY)
@@ -168,6 +176,63 @@ async def warmup(
                 except Exception:
                     logger.exception("Failed to restore STREAMING state after warmup")
 
+    async def ensure_warmup(self, width: Optional[int] = None, height: Optional[int] = None):
+        """Ensure the pipeline has been warmed up for the given resolution."""
+        if width and width > 0:
+            self.width = int(width)
+        if height and height > 0:
+            self.height = int(height)
+
+        if self._warmup_completed and self._last_warmup_resolution:
+            if (self.width, self.height) != self._last_warmup_resolution:
+                self._warmup_completed = False
+
+        if self._warmup_completed:
+            return
+
+        if not self.state_manager.is_initialized():
+            logger.debug("Skipping warmup scheduling - pipeline not initialized")
+            return
+
+        async with self._warmup_lock:
+            if self._warmup_completed:
+                return
+            if not self.state_manager.is_initialized():
+                return
+            if self._warmup_task and not self._warmup_task.done():
+                return
+
+            logger.info("Scheduling pipeline warmup for %sx%s", self.width, self.height)
+            self.disable_ingest()
+            self._warmup_task = asyncio.create_task(self._run_background_warmup())
+
+    async def _run_background_warmup(self):
+        try:
+            await self.warmup()
+        except asyncio.CancelledError:
+            logger.debug("Pipeline warmup task cancelled")
+            raise
+        except Exception:
+            logger.exception("Pipeline warmup failed")
+        finally:
+            self.enable_ingest()
+            self._warmup_task = None
+
+    async def _reset_warmup_state(self):
+        """Reset warmup bookkeeping and cancel any in-flight warmup tasks."""
+        async with self._warmup_lock:
+            if self._warmup_task and not self._warmup_task.done():
+                self._warmup_task.cancel()
+                try:
+                    await self._warmup_task
+                except asyncio.CancelledError:
+                    pass
+                except Exception:
+                    logger.debug("Warmup task raised during cancellation", exc_info=True)
+            self._warmup_task = None
+            self._warmup_completed = False
+            self._last_warmup_resolution = None
+
     async def _run_warmup(
         self,
         *,
@@ -266,6 +331,8 @@ async def set_prompts(
             skip_warmup: Skip automatic warmup even if auto_warmup is enabled
         """
         try:
+            await self._reset_warmup_state()
+
             prompt_list = prompts if isinstance(prompts, list) else [prompts]
             await self.client.set_prompts(prompt_list)
 
@@ -312,6 +379,8 @@ async def update_prompts(
             if was_streaming and should_warmup:
                 await self.state_manager.transition_to(PipelineState.READY)
 
+            await self._reset_warmup_state()
+
             prompt_list = prompts if isinstance(prompts, list) else [prompts]
             await self.client.update_prompts(prompt_list)
 
@@ -775,6 +844,7 @@ async def cleanup(self):
         # Clear cached modalities and I/O capabilities since we're resetting
         self._cached_modalities = None
         self._cached_io_capabilities = None
+        await self._reset_warmup_state()
 
         # Clear pipeline queues
         await self._clear_pipeline_queues()