better app structure

app.py

@@ -1,17 +0,0 @@
-from fastapi import FastAPI
-
-from config import args
-from device import device, torch_dtype
-from app_init import init_app
-from user_queue import user_data
-from util import get_pipeline_class
-
-print("DEVICE:", device)
-print("TORCH_DTYPE:", torch_dtype)
-args.pretty_print()
-
-app = FastAPI()
-
-pipeline_class = get_pipeline_class(args.pipeline)
-pipeline = pipeline_class(args, device, torch_dtype)
-init_app(app, user_data, args, pipeline)

app_init.py

@@ -1,163 +0,0 @@
-from fastapi import FastAPI, WebSocket, HTTPException, WebSocketDisconnect
-from fastapi.responses import StreamingResponse, JSONResponse
-from fastapi.middleware.cors import CORSMiddleware
-from fastapi.staticfiles import StaticFiles
-from fastapi import Request
-import markdown2
-
-import logging
-import traceback
-from config import Args
-from user_queue import UserData
-import uuid
-import time
-from types import SimpleNamespace
-from util import pil_to_frame, bytes_to_pil, is_firefox
-import asyncio
-import os
-import time
-
-THROTTLE = 1.0 / 120
-
-
-def init_app(app: FastAPI, user_data: UserData, args: Args, pipeline):
-    app.add_middleware(
-        CORSMiddleware,
-        allow_origins=["*"],
-        allow_credentials=True,
-        allow_methods=["*"],
-        allow_headers=["*"],
-    )
-
-    @app.websocket("/api/ws")
-    async def websocket_endpoint(websocket: WebSocket):
-        await websocket.accept()
-        user_count = user_data.get_user_count()
-        if args.max_queue_size > 0 and user_count >= args.max_queue_size:
-            print("Server is full")
-            await websocket.send_json({"status": "error", "message": "Server is full"})
-            await websocket.close()
-            return
-        try:
-            user_id = uuid.uuid4()
-            print(f"New user connected: {user_id}")
-            await user_data.create_user(user_id, websocket)
-            await websocket.send_json(
-                {"status": "connected", "message": "Connected", "userId": str(user_id)}
-            )
-            await websocket.send_json({"status": "send_frame"})
-            await handle_websocket_data(user_id, websocket)
-        except WebSocketDisconnect as e:
-            logging.error(f"WebSocket Error: {e}, {user_id}")
-            traceback.print_exc()
-        finally:
-            print(f"User disconnected: {user_id}")
-            user_data.delete_user(user_id)
-
-    async def handle_websocket_data(user_id: uuid.UUID, websocket: WebSocket):
-        if not user_data.check_user(user_id):
-            return HTTPException(status_code=404, detail="User not found")
-        last_time = time.time()
-        try:
-            while True:
-                if args.timeout > 0 and time.time() - last_time > args.timeout:
-                    await websocket.send_json(
-                        {
-                            "status": "timeout",
-                            "message": "Your session has ended",
-                            "userId": str(user_id),
-                        }
-                    )
-                    await websocket.close()
-                    return
-                data = await websocket.receive_json()
-                if data["status"] != "next_frame":
-                    asyncio.sleep(THROTTLE)
-                    continue
-
-                params = await websocket.receive_json()
-                params = pipeline.InputParams(**params)
-                info = pipeline.Info()
-                params = SimpleNamespace(**params.dict())
-                if info.input_mode == "image":
-                    image_data = await websocket.receive_bytes()
-                    if len(image_data) == 0:
-                        await websocket.send_json({"status": "send_frame"})
-                        await asyncio.sleep(THROTTLE)
-                        continue
-                    params.image = bytes_to_pil(image_data)
-                await user_data.update_data(user_id, params)
-                await websocket.send_json({"status": "wait"})
-
-        except Exception as e:
-            logging.error(f"Error: {e}")
-            traceback.print_exc()
-
-    @app.get("/api/queue")
-    async def get_queue_size():
-        queue_size = user_data.get_user_count()
-        return JSONResponse({"queue_size": queue_size})
-
-    @app.get("/api/stream/{user_id}")
-    async def stream(user_id: uuid.UUID, request: Request):
-        try:
-
-            async def generate():
-                websocket = user_data.get_websocket(user_id)
-                last_params = SimpleNamespace()
-                while True:
-                    last_time = time.time()
-                    params = await user_data.get_latest_data(user_id)
-                    if not vars(params) or params.__dict__ == last_params.__dict__:
-                        await websocket.send_json({"status": "send_frame"})
-                        await asyncio.sleep(THROTTLE)
-                        continue
-
-                    last_params = params
-                    image = pipeline.predict(params)
-
-                    if image is None:
-                        await websocket.send_json({"status": "send_frame"})
-                        await asyncio.sleep(THROTTLE)
-                        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 websocket.send_json({"status": "send_frame"})
-                    if args.debug:
-                        print(f"Time taken: {time.time() - last_time}")
-
-            return StreamingResponse(
-                generate(),
-                media_type="multipart/x-mixed-replace;boundary=frame",
-                headers={"Cache-Control": "no-cache"},
-            )
-        except Exception as e:
-            logging.error(f"Streaming Error: {e}, {user_id} ")
-            traceback.print_exc()
-            return HTTPException(status_code=404, detail="User not found")
-
-    # route to setup frontend
-    @app.get("/api/settings")
-    async def settings():
-        info_schema = pipeline.Info.schema()
-        info = pipeline.Info()
-        if info.page_content:
-            page_content = markdown2.markdown(info.page_content)
-
-        input_params = pipeline.InputParams.schema()
-        return JSONResponse(
-            {
-                "info": info_schema,
-                "input_params": input_params,
-                "max_queue_size": args.max_queue_size,
-                "page_content": page_content if info.page_content else "",
-            }
-        )
-
-    if not os.path.exists("public"):
-        os.makedirs("public")
-
-    app.mount("/", StaticFiles(directory="public", html=True), name="public")

config.py

@@ -124,4 +124,5 @@ parser.add_argument(
 )
 parser.set_defaults(taesd=USE_TAESD)
 
-args = Args(**vars(parser.parse_args()))
+config = Args(**vars(parser.parse_args()))
+config.pretty_print()

connection_manager.py

@@ -0,0 +1,116 @@
+from typing import Dict, Union
+from uuid import UUID
+import asyncio
+from fastapi import WebSocket
+from starlette.websockets import WebSocketState
+import logging
+from types import SimpleNamespace
+
+Connections = Dict[UUID, Dict[str, Union[WebSocket, asyncio.Queue]]]
+
+
+class ServerFullException(Exception):
+    """Exception raised when the server is full."""
+
+    pass
+
+
+class ConnectionManager:
+    def __init__(self):
+        self.active_connections: Connections = {}
+
+    async def connect(
+        self, user_id: UUID, websocket: WebSocket, max_queue_size: int = 0
+    ):
+        await websocket.accept()
+        user_count = self.get_user_count()
+        print(f"User count: {user_count}")
+        if max_queue_size > 0 and user_count >= max_queue_size:
+            print("Server is full")
+            await websocket.send_json({"status": "error", "message": "Server is full"})
+            await websocket.close()
+            raise ServerFullException("Server is full")
+        print(f"New user connected: {user_id}")
+        self.active_connections[user_id] = {
+            "websocket": websocket,
+            "queue": asyncio.Queue(),
+        }
+        await websocket.send_json(
+            {"status": "connected", "message": "Connected"},
+        )
+        await websocket.send_json({"status": "wait"})
+        await websocket.send_json({"status": "send_frame"})
+
+    def check_user(self, user_id: UUID) -> bool:
+        return user_id in self.active_connections
+
+    async def update_data(self, user_id: UUID, new_data: SimpleNamespace):
+        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:
+        user_session = self.active_connections.get(user_id)
+        if user_session:
+            queue = user_session["queue"]
+            try:
+                return await queue.get()
+            except asyncio.QueueEmpty:
+                return None
+
+    def delete_user(self, user_id: UUID):
+        user_session = self.active_connections.pop(user_id, None)
+        if user_session:
+            queue = user_session["queue"]
+            while not queue.empty():
+                try:
+                    queue.get_nowait()
+                except asyncio.QueueEmpty:
+                    continue
+
+    def get_user_count(self) -> int:
+        return len(self.active_connections)
+
+    def get_websocket(self, user_id: UUID) -> WebSocket:
+        user_session = self.active_connections.get(user_id)
+        if user_session:
+            websocket = user_session["websocket"]
+            if websocket.client_state == WebSocketState.CONNECTED:
+                return user_session["websocket"]
+        return None
+
+    async def disconnect(self, user_id: UUID):
+        websocket = self.get_websocket(user_id)
+        if websocket:
+            await websocket.close()
+        self.delete_user(user_id)
+
+    async def send_json(self, user_id: UUID, data: Dict):
+        try:
+            websocket = self.get_websocket(user_id)
+            if websocket:
+                await websocket.send_json(data)
+        except Exception as e:
+            logging.error(f"Error: Send json: {e}")
+
+    async def receive_json(self, user_id: UUID) -> Dict:
+        try:
+            websocket = self.get_websocket(user_id)
+            if websocket:
+                return await websocket.receive_json()
+        except Exception as e:
+            logging.error(f"Error: Receive json: {e}")
+
+    async def receive_bytes(self, user_id: UUID) -> bytes:
+        try:
+            websocket = self.get_websocket(user_id)
+            if websocket:
+                return await websocket.receive_bytes()
+        except Exception as e:
+            logging.error(f"Error: Receive bytes: {e}")

frontend/src/lib/lcmLive.ts

@@ -6,6 +6,7 @@ export enum LCMLiveStatus {
     DISCONNECTED = "disconnected",
     WAIT = "wait",
     SEND_FRAME = "send_frame",
+    TIMEOUT = "timeout",
 }
 
 const initStatus: LCMLiveStatus = LCMLiveStatus.DISCONNECTED;
@@ -19,8 +20,9 @@ export const lcmLiveActions = {
         return new Promise((resolve, reject) => {
 
             try {
+                const userId = crypto.randomUUID();
                 const websocketURL = `${window.location.protocol === "https:" ? "wss" : "ws"
-                    }:${window.location.host}/api/ws`;
+                    }:${window.location.host}/api/ws/${userId}`;
 
                 websocket = new WebSocket(websocketURL);
                 websocket.onopen = () => {
@@ -37,9 +39,9 @@ export const lcmLiveActions = {
                     const data = JSON.parse(event.data);
                     switch (data.status) {
                         case "connected":
-                            const userId = data.userId;
                             lcmLiveStatus.set(LCMLiveStatus.CONNECTED);
                             streamId.set(userId);
+                            resolve({ status: "connected", userId });
                             break;
                         case "send_frame":
                             lcmLiveStatus.set(LCMLiveStatus.SEND_FRAME);
@@ -54,14 +56,16 @@ export const lcmLiveActions = {
                             break;
                         case "timeout":
                             console.log("timeout");
-                            lcmLiveStatus.set(LCMLiveStatus.DISCONNECTED);
+                            lcmLiveStatus.set(LCMLiveStatus.TIMEOUT);
                             streamId.set(null);
-                            resolve({ status: "timeout" });
+                            reject(new Error("timeout"));
+                            break;
                         case "error":
                             console.log(data.message);
                             lcmLiveStatus.set(LCMLiveStatus.DISCONNECTED);
                             streamId.set(null);
                             reject(new Error(data.message));
+                            break;
                     }
                 };
 
@@ -85,12 +89,11 @@ export const lcmLiveActions = {
         }
     },
     async stop() {
-
+        lcmLiveStatus.set(LCMLiveStatus.DISCONNECTED);
         if (websocket) {
             websocket.close();
         }
         websocket = null;
-        lcmLiveStatus.set(LCMLiveStatus.DISCONNECTED);
         streamId.set(null);
     },
 };

frontend/src/routes/+page.svelte

@@ -20,7 +20,6 @@
   let currentQueueSize: number = 0;
   let queueCheckerRunning: boolean = false;
   let warningMessage: string = '';
-
   onMount(() => {
     getSettings();
   });
@@ -59,7 +58,9 @@
   }
 
   $: isLCMRunning = $lcmLiveStatus !== LCMLiveStatus.DISCONNECTED;
-
+  $: if ($lcmLiveStatus === LCMLiveStatus.TIMEOUT) {
+    warningMessage = 'Session timed out. Please try again.';
+  }
   let disabled = false;
   async function toggleLcmLive() {
     try {
@@ -70,7 +71,6 @@
         }
         disabled = true;
         await lcmLiveActions.start(getSreamdata);
-        warningMessage = 'Timeout, please try again.';
         disabled = false;
         toggleQueueChecker(false);
       } else {

main.py

@@ -0,0 +1,184 @@
+from fastapi import FastAPI, WebSocket, HTTPException, WebSocketDisconnect
+from fastapi.responses import StreamingResponse, JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.staticfiles import StaticFiles
+from fastapi import Request
+import markdown2
+
+import logging
+from config import config, Args
+from connection_manager import ConnectionManager
+import uuid
+import time
+from types import SimpleNamespace
+from util import pil_to_frame, bytes_to_pil, is_firefox, get_pipeline_class
+from device import device, torch_dtype
+import asyncio
+import os
+import time
+import torch
+
+
+THROTTLE = 1.0 / 120
+
+
+class App:
+    def __init__(self, config: Args, pipeline):
+        self.args = config
+        self.pipeline = pipeline
+        self.app = FastAPI()
+        self.conn_manager = ConnectionManager()
+        self.init_app()
+
+    def init_app(self):
+        self.app.add_middleware(
+            CORSMiddleware,
+            allow_origins=["*"],
+            allow_credentials=True,
+            allow_methods=["*"],
+            allow_headers=["*"],
+        )
+
+        @self.app.websocket("/api/ws/{user_id}")
+        async def websocket_endpoint(user_id: uuid.UUID, websocket: WebSocket):
+            try:
+                await self.conn_manager.connect(
+                    user_id, websocket, self.args.max_queue_size
+                )
+                await handle_websocket_data(user_id)
+            except ServerFullException as e:
+                logging.error(f"Server Full: {e}")
+            finally:
+                await self.conn_manager.disconnect(user_id)
+                logging.info(f"User disconnected: {user_id}")
+
+        async def handle_websocket_data(user_id: uuid.UUID):
+            if not self.conn_manager.check_user(user_id):
+                return HTTPException(status_code=404, detail="User not found")
+            last_time = time.time()
+            try:
+                while True:
+                    if (
+                        self.args.timeout > 0
+                        and time.time() - last_time > self.args.timeout
+                    ):
+                        await self.conn_manager.send_json(
+                            user_id,
+                            {
+                                "status": "timeout",
+                                "message": "Your session has ended",
+                            },
+                        )
+                        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"})
+
+            except Exception as e:
+                logging.error(f"Websocket Error: {e}, {user_id} ")
+                await self.conn_manager.disconnect(user_id)
+
+        @self.app.get("/api/queue")
+        async def get_queue_size():
+            queue_size = self.conn_manager.get_user_count()
+            return JSONResponse({"queue_size": queue_size})
+
+        @self.app.get("/api/stream/{user_id}")
+        async def stream(user_id: uuid.UUID, request: Request):
+            try:
+
+                async def generate():
+                    last_params = SimpleNamespace()
+                    while True:
+                        last_time = time.time()
+                        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"}
+                            )
+                            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}")
+
+                return StreamingResponse(
+                    generate(),
+                    media_type="multipart/x-mixed-replace;boundary=frame",
+                    headers={"Cache-Control": "no-cache"},
+                )
+            except Exception as e:
+                logging.error(f"Streaming Error: {e}, {user_id} ")
+                return HTTPException(status_code=404, detail="User not found")
+
+        # route to setup frontend
+        @self.app.get("/api/settings")
+        async def settings():
+            info_schema = pipeline.Info.schema()
+            info = pipeline.Info()
+            if info.page_content:
+                page_content = markdown2.markdown(info.page_content)
+
+            input_params = pipeline.InputParams.schema()
+            return JSONResponse(
+                {
+                    "info": info_schema,
+                    "input_params": input_params,
+                    "max_queue_size": self.args.max_queue_size,
+                    "page_content": page_content if info.page_content else "",
+                }
+            )
+
+        if not os.path.exists("public"):
+            os.makedirs("public")
+
+        self.app.mount("/", StaticFiles(directory="public", html=True), name="public")
+
+
+pipeline_class = get_pipeline_class(config.pipeline)
+pipeline = pipeline_class(config, device, torch_dtype)
+app = App(config, pipeline).app
+
+if __name__ == "__main__":
+    import uvicorn
+
+    uvicorn.run(
+        "main:app",
+        host=config.host,
+        port=config.port,
+        reload=config.reload,
+        ssl_certfile=config.ssl_certfile,
+        ssl_keyfile=config.ssl_keyfile,
+    )

run.py

@@ -1,12 +0,0 @@
-if __name__ == "__main__":
-    import uvicorn
-    from config import args
-
-    uvicorn.run(
-        "app:app",
-        host=args.host,
-        port=args.port,
-        reload=args.reload,
-        ssl_certfile=args.ssl_certfile,
-        ssl_keyfile=args.ssl_keyfile,
-    )

user_queue.py

@@ -1,63 +0,0 @@
-from typing import Dict
-from uuid import UUID
-import asyncio
-from fastapi import WebSocket
-from types import SimpleNamespace
-from typing import Dict
-from typing import Union
-
-UserDataContent = Dict[UUID, Dict[str, Union[WebSocket, asyncio.Queue]]]
-
-
-class UserData:
-    def __init__(self):
-        self.data_content: Dict[UUID, UserDataContent] = {}
-
-    async def create_user(self, user_id: UUID, websocket: WebSocket):
-        self.data_content[user_id] = {
-            "websocket": websocket,
-            "queue": asyncio.Queue(),
-        }
-        await asyncio.sleep(1)
-
-    def check_user(self, user_id: UUID) -> bool:
-        return user_id in self.data_content
-
-    async def update_data(self, user_id: UUID, new_data: SimpleNamespace):
-        user_session = self.data_content[user_id]
-        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:
-        user_session = self.data_content[user_id]
-        queue = user_session["queue"]
-
-        try:
-            return await queue.get()
-        except asyncio.QueueEmpty:
-            return None
-
-    def delete_user(self, user_id: UUID):
-        user_session = self.data_content[user_id]
-        queue = user_session["queue"]
-        while not queue.empty():
-            try:
-                queue.get_nowait()
-            except asyncio.QueueEmpty:
-                continue
-        if user_id in self.data_content:
-            del self.data_content[user_id]
-
-    def get_user_count(self) -> int:
-        return len(self.data_content)
-
-    def get_websocket(self, user_id: UUID) -> WebSocket:
-        return self.data_content[user_id]["websocket"]
-
-
-user_data = UserData()

util.py

@@ -1,7 +1,5 @@
 from importlib import import_module
 from types import ModuleType
-from typing import Dict, Any
-from pydantic import BaseModel as PydanticBaseModel, Field
 from PIL import Image
 import io