feat: add cleanup

Author: varshith15

src/comfystream/client.py

@@ -15,7 +15,6 @@
 class ComfyStreamClient:
     def __init__(self, max_workers: int = 1, **kwargs):
         config = Configuration(**kwargs)
-        # TODO: Need to handle cleanup for EmbeddedComfyClient if not using async context manager?
         self.comfy_client = EmbeddedComfyClient(config, max_workers=max_workers)
         self.running_prompts = {} # To be used for cancelling tasks
         self.current_prompts = []
@@ -43,6 +42,13 @@ async def run_prompt(self, prompt_index: int):
                 logger.error(f"Error type: {type(e)}")
                 raise
 
+    async def cleanup(self):
+        for task in self.running_prompts.values():
+            await task.cancel()
+
+        if self.comfy_client.is_running:
+            await self.comfy_client.__aexit__()
+
     def put_video_input(self, frame):
         if tensor_cache.image_inputs.full():
             tensor_cache.image_inputs.get(block=True)