Fixed bug in accelerator.

accelerator.py

@@ -1,5 +1,5 @@
 from typing import Union
-from fastapi.websockets import WebSocket, WebSocketState
+from fastapi.websockets import WebSocket, WebSocketDisconnect
 
 class Accelerator:
   ws: Union[WebSocket, None]
@@ -10,11 +10,9 @@ class Accelerator:
     await ws.accept()
     self.ws = ws
 
-  async def accelerate(self, input):
-    while ws.client_state == WebSocketState.CONNECTING: pass
-    
-    if ws.client_state == WebSocketState.CONNECTED:
+  async def accelerate(self, input):    
+    try:
       await self.ws.send_text(input)
       return await self.ws.receive_text()
-    elif ws.client_state == WebSocketState.DISCONNECTED:
-      ws = None
+    except WebSocketDisconnect:
+      self.ws = None