Add live stream memory cleanup hook

Author: fastagents

live/streamdiffusion/pipeline/pipeline.py

@@ -2,6 +2,8 @@
 import logging
 import asyncio
 import base64
+import ctypes
+import gc
 import re
 from pathlib import Path
 from typing import Dict, List, Optional, Any, cast
@@ -30,6 +32,10 @@
 ENGINES_DIR = Path("./engines")
 # this one is used only for realesrgan_trt which has ./models hardcoded
 LOCAL_MODELS_DIR = Path("./models")
+STREAM_CLEANUP_GC_ENV = "LIVE_PIPELINE_STREAM_CLEANUP_GC"
+STREAM_CLEANUP_MALLOC_TRIM_ENV = "LIVE_PIPELINE_STREAM_CLEANUP_MALLOC_TRIM"
+STREAM_CLEANUP_CUDA_CACHE_ENV = "LIVE_PIPELINE_STREAM_CLEANUP_CUDA_CACHE"
+STREAM_CLEANUP_STYLE_CACHE_ENV = "LIVE_PIPELINE_STREAM_CLEANUP_STYLE_CACHE"
 
 class StreamDiffusion(Pipeline):
     def __init__(self):
@@ -280,13 +286,51 @@ async def stop(self):
             self.params = None
             self.frame_queue = asyncio.Queue()
 
+    async def cleanup_stream(self):
+        async with self._pipeline_lock:
+            _clear_asyncio_queue(self.frame_queue)
+            if _env_enabled(STREAM_CLEANUP_STYLE_CACHE_ENV, default=False):
+                self._cached_style_image_tensor = None
+                self._cached_style_image_url = None
+
+        if _env_enabled(STREAM_CLEANUP_GC_ENV, default=True):
+            await asyncio.to_thread(gc.collect)
+        if _env_enabled(STREAM_CLEANUP_MALLOC_TRIM_ENV, default=True):
+            await asyncio.to_thread(_malloc_trim)
+        if _env_enabled(STREAM_CLEANUP_CUDA_CACHE_ENV, default=False) and torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        logging.info("StreamDiffusion stream cleanup complete")
+
     @classmethod
     def prepare_models(cls):
         from .prepare import prepare_streamdiffusion_models
 
         prepare_streamdiffusion_models()
 
 
+def _env_enabled(name: str, *, default: bool) -> bool:
+    raw = os.getenv(name)
+    if raw is None:
+        return default
+    return raw.strip().lower() in {"1", "true", "yes", "on"}
+
+
+def _malloc_trim() -> None:
+    try:
+        libc = ctypes.CDLL("libc.so.6")
+        libc.malloc_trim(0)
+    except Exception:
+        logging.debug("malloc_trim unavailable", exc_info=True)
+
+
+def _clear_asyncio_queue(queue: asyncio.Queue) -> None:
+    while True:
+        try:
+            queue.get_nowait()
+        except asyncio.QueueEmpty:
+            return
+
+
 def _prepare_controlnet_configs(params: StreamDiffusionParams) -> Optional[List[Dict[str, Any]]]:
     """Prepare ControlNet configurations for wrapper"""
     if not params.controlnets:

runner/src/runner/live/pipelines/interface.py

@@ -108,6 +108,15 @@ async def stop(self):
         """
         pass
 
+    async def cleanup_stream(self):
+        """Release per-stream state while keeping the warmed pipeline alive.
+
+        Called after a stream goes offline. Implementations should clear
+        frame queues, per-stream caches, and allocator state that does not need
+        to survive into the next stream, without unloading model weights.
+        """
+        pass
+
     @classmethod
     @abstractmethod
     def prepare_models(cls):

runner/src/runner/live/process/process.py

@@ -28,6 +28,8 @@
 
 from .loading_overlay import LoadingOverlayRenderer
 
+STREAM_CLEANUP_CONTROL_KEY = "__stream_cleanup__"
+
 
 class PipelineProcess:
     @staticmethod
@@ -139,20 +141,28 @@ def _set_pipeline_ready(self, ready: bool):
     def update_params(self, params: dict):
         self.param_update_queue.put(params)
 
-    def reset_stream(self, request_id: str, manifest_id: str, stream_id: str):
+    def reset_stream(
+        self,
+        request_id: str,
+        manifest_id: str,
+        stream_id: str,
+        *,
+        cleanup_stream: bool = False,
+    ):
         # Clear queues to avoid using frames from previous sessions
         clear_queue(self.input_queue)
         clear_queue(self.output_queue)
         clear_queue(self.param_update_queue)
         clear_queue(self.error_queue)
         clear_queue(self.log_queue)
-        self.param_update_queue.put(
-            {
-                "request_id": request_id,
-                "manifest_id": manifest_id,
-                "stream_id": stream_id,
-            }
-        )
+        message = {
+            "request_id": request_id,
+            "manifest_id": manifest_id,
+            "stream_id": stream_id,
+        }
+        if cleanup_stream:
+            message[STREAM_CLEANUP_CONTROL_KEY] = True
+        self.param_update_queue.put(message)
 
     # TODO: Once audio is implemented, combined send_input with input_loop
     # We don't need additional queueing as comfystream already maintains a queue
@@ -330,7 +340,7 @@ async def _param_update_loop(self, pipeline: Pipeline, overlay: LoadingOverlayRe
         while not self.is_done():
             reload_task = None
             try:
-                params = await self._get_latest_params(timeout=0.1)
+                params = await self._get_latest_params(timeout=0.1, pipeline=pipeline, overlay=overlay)
                 if params is None:
                     continue
 
@@ -382,7 +392,12 @@ async def _param_update_loop(self, pipeline: Pipeline, overlay: LoadingOverlayRe
                 except Exception:
                     logging.warning("Failed to prewarm loading overlay caches", exc_info=True)
 
-    async def _get_latest_params(self, timeout: float) -> dict | None:
+    async def _get_latest_params(
+        self,
+        timeout: float,
+        pipeline: Pipeline,
+        overlay: LoadingOverlayRenderer,
+    ) -> dict | None:
         """
         Get the latest params from the param_update_queue, skipping stale entries before the latest. Already filters
         and processes params that are only logging updates. Waits for timeout seconds for a new params entry, or returns
@@ -396,20 +411,44 @@ async def _get_latest_params(self, timeout: float) -> dict | None:
         except queue.Empty:
             return None
 
-        if self._handle_logging_params(params):
+        if await self._handle_process_message(params, pipeline, overlay):
             params = None
 
         # Drain the params queue to get the latest params
         while not self.param_update_queue.empty():
             try:
                 new_params = self.param_update_queue.get_nowait()
-                if not self._handle_logging_params(new_params):
+                if not await self._handle_process_message(new_params, pipeline, overlay):
                     params = new_params
             except queue.Empty:
                 break
 
         return params
 
+    async def _handle_process_message(
+        self,
+        params: dict,
+        pipeline: Pipeline,
+        overlay: LoadingOverlayRenderer,
+    ) -> bool:
+        consumed = self._handle_logging_params(params)
+        if isinstance(params, dict) and params.get(STREAM_CLEANUP_CONTROL_KEY):
+            await self._cleanup_stream_state(pipeline, overlay)
+            consumed = True
+        return consumed
+
+    async def _cleanup_stream_state(self, pipeline: Pipeline, overlay: LoadingOverlayRenderer):
+        logging.info("PipelineProcess: stream cleanup requested")
+        clear_queue(self.input_queue)
+        clear_queue(self.output_queue)
+        overlay.end_reload()
+        try:
+            await pipeline.cleanup_stream()
+        except Exception as e:
+            self._report_error("Error cleaning up pipeline stream state", e)
+        else:
+            logging.info("PipelineProcess: stream cleanup complete")
+
     def _report_error(self, msg: str, error: Exception | None = None, silent=False):
         if not silent:
             logging.error(msg, exc_info=error)

runner/src/runner/live/process/process_guardian.py

@@ -301,8 +301,9 @@ async def _monitor_loop(self):
                     continue
 
                 if state == PipelineState.OFFLINE:
-                    # Revert to initial params when the stream stops
-                    self.process.reset_stream("", "", "")
+                    # Revert to initial params and release stream-scoped state
+                    # while keeping the warmed pipeline alive.
+                    self.process.reset_stream("", "", "", cleanup_stream=True)
                     self.process.update_params(self.pipeline_spec.initial_params)
 
                 if state != PipelineState.ERROR:

runner/src/runner/live/streamer/protocol/trickle.py

@@ -1,5 +1,6 @@
 import asyncio
 import logging
+import os
 import queue
 import json
 from typing import AsyncGenerator, Optional
@@ -11,14 +12,49 @@
 from .protocol import StreamProtocol
 from .last_value_cache import LastValueCache
 
+DEFAULT_PUBLISH_QUEUE_MAXSIZE = 1
+
+
+def _publish_queue_maxsize() -> int:
+    raw = os.environ.get("LIVE_TRICKLE_PUBLISH_QUEUE_MAXSIZE")
+    if raw is None or raw.strip() == "":
+        return DEFAULT_PUBLISH_QUEUE_MAXSIZE
+    try:
+        value = int(raw)
+    except ValueError:
+        logging.warning(
+            "Invalid LIVE_TRICKLE_PUBLISH_QUEUE_MAXSIZE=%r, using %s",
+            raw,
+            DEFAULT_PUBLISH_QUEUE_MAXSIZE,
+        )
+        return DEFAULT_PUBLISH_QUEUE_MAXSIZE
+    return max(1, value)
+
+
+def _put_drop_oldest(_queue: queue.Queue, item) -> int:
+    dropped = 0
+    while True:
+        try:
+            _queue.put_nowait(item)
+            return dropped
+        except queue.Full:
+            try:
+                old = _queue.get_nowait()
+            except queue.Empty:
+                continue
+            dropped += 1
+            del old
+
+
 class TrickleProtocol(StreamProtocol):
     def __init__(self, subscribe_url: str, publish_url: str, control_url: Optional[str] = None, events_url: Optional[str] = None, input_width: Optional[int] = DEFAULT_WIDTH, input_height: Optional[int] = DEFAULT_HEIGHT, output_width: Optional[int] = None, output_height: Optional[int] = None):
         self.subscribe_url = subscribe_url
         self.publish_url = publish_url
         self.control_url = control_url
         self.events_url = events_url
         self.subscribe_queue = queue.Queue[InputFrame]()
-        self.publish_queue = queue.Queue[OutputFrame]()
+        self.publish_queue = queue.Queue[OutputFrame](maxsize=_publish_queue_maxsize())
+        self.publish_queue_dropped = 0
         self.control_subscriber = None
         self.events_publisher = None
         self.subscribe_task = None
@@ -30,7 +66,8 @@ def __init__(self, subscribe_url: str, publish_url: str, control_url: Optional[s
 
     async def start(self):
         self.subscribe_queue = queue.Queue[InputFrame]()
-        self.publish_queue = queue.Queue[OutputFrame]()
+        self.publish_queue = queue.Queue[OutputFrame](maxsize=_publish_queue_maxsize())
+        self.publish_queue_dropped = 0
         metadata_cache = LastValueCache[dict]() # to pass video metadata from decoder to encoder
         self.subscribe_task = asyncio.create_task(
             media.run_subscribe(self.subscribe_url, self.subscribe_queue.put, metadata_cache.put, self.emit_monitoring_event, self.input_width, self.input_height)
@@ -49,7 +86,7 @@ async def stop(self):
 
         # send sentinel None values to stop the trickle tasks gracefully
         self.subscribe_queue.put(None)
-        self.publish_queue.put(None)
+        _put_drop_oldest(self.publish_queue, None)
 
         if self.control_subscriber:
             await self.control_subscriber.close()
@@ -86,17 +123,30 @@ def dequeue_frame():
             # TEMP: Put audio immediately into the publish queue
             # TODO: Remove once there is ComfyUI audio support
             if isinstance(image, AudioFrame):
-                publish_queue.put(AudioOutput([image]))
+                self.publish_queue_dropped += _put_drop_oldest(
+                    publish_queue,
+                    AudioOutput([image]),
+                )
                 continue
             yield image
 
     async def egress_loop(self, output_frames: AsyncGenerator[OutputFrame, None]):
         publish_queue = self.publish_queue
         def enqueue_bytes(frame: OutputFrame):
-            publish_queue.put(frame)
+            return _put_drop_oldest(publish_queue, frame)
 
         async for frame in output_frames:
-            await asyncio.to_thread(enqueue_bytes, frame)
+            try:
+                dropped = await asyncio.to_thread(enqueue_bytes, frame)
+                self.publish_queue_dropped += dropped
+                if dropped and self.publish_queue_dropped % 100 == 1:
+                    logging.warning(
+                        "Dropped stale publish frames dropped_total=%s queue_maxsize=%s",
+                        self.publish_queue_dropped,
+                        publish_queue.maxsize,
+                    )
+            finally:
+                del frame
 
     async def emit_monitoring_event(self, event: dict, queue_event_type: str = "ai_stream_events"):
         if not self.events_publisher:
@@ -134,4 +184,3 @@ async def control_loop(self, done: asyncio.Event) -> AsyncGenerator[dict, None]:
             except Exception:
                 logging.error(f"Error in control loop", exc_info=True)
                 continue
-