better Websocket handling

connection_manager.py

@@ -48,11 +48,6 @@ class ConnectionManager:
         user_session = self.active_connections.get(user_id)
         if user_session:
             queue = user_session["queue"]
-            while not queue.empty():
-                try:
-                    queue.get_nowait()
-                except asyncio.QueueEmpty:
-                    continue
             await queue.put(new_data)
 
     async def get_latest_data(self, user_id: UUID) -> SimpleNamespace:

main.py

@@ -7,7 +7,7 @@ import markdown2
 
 import logging
 from config import config, Args
-from connection_manager import ConnectionManager
+from connection_manager import ConnectionManager, ServerFullException
 import uuid
 import time
 from types import SimpleNamespace
@@ -72,25 +72,22 @@ class App:
                         await self.conn_manager.disconnect(user_id)
                         return
                     data = await self.conn_manager.receive_json(user_id)
-                    if data["status"] != "next_frame":
-                        asyncio.sleep(THROTTLE)
-                        continue
-
-                    params = await self.conn_manager.receive_json(user_id)
-                    params = pipeline.InputParams(**params)
-                    info = pipeline.Info()
-                    params = SimpleNamespace(**params.dict())
-                    if info.input_mode == "image":
-                        image_data = await self.conn_manager.receive_bytes(user_id)
-                        if len(image_data) == 0:
-                            await self.conn_manager.send_json(
-                                user_id, {"status": "send_frame"}
-                            )
-                            await asyncio.sleep(THROTTLE)
-                            continue
-                        params.image = bytes_to_pil(image_data)
-                    await self.conn_manager.update_data(user_id, params)
-                    await self.conn_manager.send_json(user_id, {"status": "wait"})
+                    if data["status"] == "next_frame":
+                        info = pipeline.Info()
+                        params = await self.conn_manager.receive_json(user_id)
+                        params = pipeline.InputParams(**params)
+                        params = SimpleNamespace(**params.dict())
+                        if info.input_mode == "image":
+                            image_data = await self.conn_manager.receive_bytes(user_id)
+                            if len(image_data) == 0:
+                                await self.conn_manager.send_json(
+                                    user_id, {"status": "send_frame"}
+                                )
+                                continue
+                            params.image = bytes_to_pil(image_data)
+
+                        await self.conn_manager.update_data(user_id, params)
+                        await self.conn_manager.send_json(user_id, {"status": "wait"})
 
             except Exception as e:
                 logging.error(f"Websocket Error: {e}, {user_id} ")
@@ -109,28 +106,22 @@ class App:
                     last_params = SimpleNamespace()
                     while True:
                         last_time = time.time()
+                        await self.conn_manager.send_json(
+                            user_id, {"status": "send_frame"}
+                        )
                         params = await self.conn_manager.get_latest_data(user_id)
-                        if not vars(params) or params.__dict__ == last_params.__dict__:
-                            await self.conn_manager.send_json(
-                                user_id, {"status": "send_frame"}
-                            )
+                        if params.__dict__ == last_params.__dict__ or params is None:
+                            await asyncio.sleep(THROTTLE)
                             continue
-
                         last_params = params
                         image = pipeline.predict(params)
                         if image is None:
-                            await self.conn_manager.send_json(
-                                user_id, {"status": "send_frame"}
-                            )
                             continue
                         frame = pil_to_frame(image)
                         yield frame
                         # https://bugs.chromium.org/p/chromium/issues/detail?id=1250396
                         if not is_firefox(request.headers["user-agent"]):
                             yield frame
-                        await self.conn_manager.send_json(
-                            user_id, {"status": "send_frame"}
-                        )
                         if self.args.debug:
                             print(f"Time taken: {time.time() - last_time}")
 

pipelines/txt2img.py

@@ -7,10 +7,10 @@ try:
 except:
     pass
 
-import psutil
 from config import Args
 from pydantic import BaseModel, Field
 from PIL import Image
+from typing import List
 
 base_model = "SimianLuo/LCM_Dreamshaper_v7"
 taesd_model = "madebyollin/taesd"