feat: comfyui stop (#432)

* feat: comfyui stop

* fix: pipeline stop

* fix: nit

* fix: process restart

* fix: add todo

runner/app/live/pipelines/comfyui.py

@@ -297,3 +297,8 @@ def update_params(self, **params):
         logging.info(f"ComfyUI Pipeline Prompt: {new_params.prompt}")
         self.client.set_prompt(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):
+        if self.client.comfy_client.is_running:
+            await self.client.comfy_client.__aexit__(None, None, None)

runner/app/live/pipelines/interface.py

@@ -46,3 +46,10 @@ def update_params(self, **params):
             **params: Implementation-specific parameters
         """
         pass
+
+    async def stop(self):
+        """Stop the pipeline.
+
+        Called once when the pipeline is no longer needed.
+        """
+        pass

runner/app/live/streamer/process.py

@@ -37,10 +37,13 @@ def __init__(self, pipeline_name: str):
         self.start_time = 0.0
 
     async def stop(self):
-        await asyncio.to_thread(self._stop_sync)
+        #TODO: syncrounous call might be blocking, figure out a better way
+        self._stop_sync()
 
     def _stop_sync(self):
         self.done.set()
+        time.sleep(2) # wait for the pipeline to stop
+
         if not self.process.is_alive():
             logging.info("Process already not alive")
             return
@@ -100,6 +103,7 @@ def get_recent_logs(self, n=10) -> list[str]:
 
     def process_loop(self):
         self._setup_logging()
+        pipeline = None
 
         def report_error(error_msg: str):
             error_event = {
@@ -170,6 +174,15 @@ def report_error(error_msg: str):
                     report_error(f"Error processing frame: {e}")
         except Exception as e:
             report_error(f"Error in process run method: {e}")
+        finally:
+            self._cleanup_pipeline(pipeline)
+
+    def _cleanup_pipeline(self, pipeline):
+        if pipeline is not None:
+            try:
+                asyncio.get_event_loop().run_until_complete(pipeline.stop())
+            except Exception as e:
+                logging.error(f"Error stopping pipeline: {e}")
 
     def _setup_logging(self):
         level = (