fix: cleanup

Author: varshith15

server/app.py

@@ -38,8 +38,12 @@ def __init__(self, track: MediaStreamTrack, pipeline):
 
     async def collect_frames(self):
         while True:
-            frame = await self.track.recv()
-            await self.pipeline.put_video_frame(frame)
+            try:
+                frame = await self.track.recv()
+                await self.pipeline.put_video_frame(frame)
+            except Exception as e:
+                await self.pipeline.cleanup()
+                raise Exception(f"Error collecting video frames: {str(e)}")
 
     async def recv(self):
         return await self.pipeline.get_processed_video_frame()
@@ -55,8 +59,12 @@ def __init__(self, track: MediaStreamTrack, pipeline):
 
     async def collect_frames(self):
         while True:
-            frame = await self.track.recv()
-            await self.pipeline.put_audio_frame(frame)
+            try:
+                frame = await self.track.recv()
+                await self.pipeline.put_audio_frame(frame)
+            except Exception as e:
+                await self.pipeline.cleanup()
+                raise Exception(f"Error collecting audio frames: {str(e)}")
 
     async def recv(self):
         return await self.pipeline.get_processed_audio_frame()

server/pipeline.py

@@ -106,4 +106,7 @@ async def get_processed_audio_frame(self):
     async def get_nodes_info(self) -> Dict[str, Any]:
         """Get information about all nodes in the current prompt including metadata."""
         nodes_info = await self.client.get_available_nodes()
-        return nodes_info
+        return nodes_info
+    
+    async def cleanup(self):
+        await self.client.cleanup()

src/comfystream/client.py

@@ -18,6 +18,7 @@ def __init__(self, max_workers: int = 1, **kwargs):
         self.comfy_client = EmbeddedComfyClient(config, max_workers=max_workers)
         self.running_prompts = {} # To be used for cancelling tasks
         self.current_prompts = []
+        self.cleanup_lock = asyncio.Lock()
 
     async def set_prompts(self, prompts: List[PromptDictInput]):
         self.current_prompts = [convert_prompt(prompt) for prompt in prompts]
@@ -38,16 +39,39 @@ async def run_prompt(self, prompt_index: int):
             try:
                 await self.comfy_client.queue_prompt(self.current_prompts[prompt_index])
             except Exception as e:
+                await self.cleanup()
                 logger.error(f"Error running prompt: {str(e)}")
-                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__()
+        async with self.cleanup_lock:
+            for task in self.running_prompts.values():
+                task.cancel()
+                try:
+                    await task
+                except asyncio.CancelledError:
+                    pass
+            self.running_prompts.clear()
+
+            if self.comfy_client.is_running:
+                await self.comfy_client.__aexit__()
+
+            await self.cleanup_queues()
+            logger.info("Client cleanup complete")
+
+        
+    async def cleanup_queues(self):
+        while not tensor_cache.image_inputs.empty():
+            tensor_cache.image_inputs.get()
+
+        while not tensor_cache.audio_inputs.empty():
+            tensor_cache.audio_inputs.get()
+
+        while not tensor_cache.image_outputs.empty():
+            await tensor_cache.image_outputs.get()
+
+        while not tensor_cache.audio_outputs.empty():
+            await tensor_cache.audio_outputs.get()
 
     def put_video_input(self, frame):
         if tensor_cache.image_inputs.full():