Removed socketio

im-stuff · Feb 25, 2023 · 70b3311 · 70b3311
1 parent a9c5784
commit 70b3311
Show file tree

Hide file tree

Showing 3 changed files with 82 additions and 41 deletions.
diff --git a/requirements.txt b/requirements.txt
@@ -9,4 +9,3 @@ safetensors
 pytorch_lightning
 aiohttp
 accelerate
-python-socketio
diff --git a/server.py b/server.py
@@ -3,6 +3,8 @@
 import asyncio
 import nodes
 import main
+import uuid
+import json
 
 try:
     import aiohttp
@@ -14,32 +16,33 @@
     print("pip install -r requirements.txt")
     sys.exit()
 
-try:
-    import socketio
-except ImportError:
-    print("Module 'python-socketio' not installed. Please install it via:")
-    print("pip install python-socketio")
-    print("or")
-    print("pip install -r requirements.txt")
-    sys.exit()
-
-
 class PromptServer():
     def __init__(self, loop):
         self.prompt_queue = None
         self.loop = loop
         self.messages = asyncio.Queue()
         self.number = 0
         self.app = web.Application()
-        self.sio = socketio.AsyncServer()
-        self.sio.attach(self.app)
+        self.sockets = dict()
         self.web_root = os.path.join(os.path.dirname(
             os.path.realpath(__file__)), "webshit")
         routes = web.RouteTableDef()
 
-        @self.sio.event
-        async def connect(sid, environ):
-            await self.sio.emit("status", self.get_queue_info(), sid)
+        @routes.get('/ws')
+        async def websocket_handler(request):
+            ws = web.WebSocketResponse()
+            await ws.prepare(request)
+            sid = uuid.uuid4().hex
+            self.sockets[sid] = ws
+            try:
+                # Send initial state to the new client
+                await self.send("status", { "status": self.get_queue_info(), 'sid': sid }, sid)
+                async for msg in ws:
+                    if msg.type == aiohttp.WSMsgType.ERROR:
+                        print('ws connection closed with exception %s' % ws.exception())
+            finally:
+                self.sockets.pop(sid)
+            return ws
 
         @routes.get("/")
         async def get_root(request):
@@ -164,14 +167,23 @@ def get_queue_info(self):
         return prompt_info
 
     async def send(self, event, data, sid=None):
-        await self.sio.emit(event, data, to=sid)
+        message = {"type": event, "data": data}
+
+        if isinstance(message, str) == False:
+            message = json.dumps(message)
+
+        if sid is None:
+            for ws in self.sockets.values():
+                await ws.send_str(message)
+        elif sid in self.sockets:
+            await self.sockets[sid].send_str(message)
 
     def send_sync(self, event, data, sid=None):
         self.loop.call_soon_threadsafe(
             self.messages.put_nowait, (event, data, sid))
 
     def queue_updated(self):
-        self.send_sync("status", self.get_queue_info())
+        self.send_sync("status", { "status": self.get_queue_info() })
 
     async def publish_loop(self):
         while True:

diff --git a/webshit/index.html b/webshit/index.html
@@ -2,7 +2,6 @@
 <head>
 	<link rel="stylesheet" type="text/css" href="litegraph.css">
 	<script type="text/javascript" src="litegraph.core.js"></script>
-	<script type="text/javascript" src="socket.io.min.js"></script>
 </head>
 <style>
 .customtext_input {
@@ -627,46 +626,77 @@
         document.getElementById("queuesize").innerHTML = "Queue size: " + (data ? data.exec_info.queue_remaining : "ERR");
     }
 
+    //fix for colab and other things that don't support websockets.
+    function manually_fetch_queue() {
+        fetch('/prompt')
+            .then(response => response.json())
+            .then(data => {
+                updateStatus(data);
+            }).catch((response) => {updateStatus(null)});
+    }
+
+    let ws;
     function createSocket(isReconnect) {
-        let opened = false;
-        const ws = io();
+        if(ws) return;
 
-        ws.on("connect", () => {
-            clientId = ws.id;
+        let opened = false;
+        ws = new WebSocket(`ws${window.location.protocol === "https:"? "s" : ""}://${location.host}/ws`);
 
-            if(opened) {
+        ws.addEventListener("open", () => {
+            opened = true;
+            if(isReconnect) {
                 closeModal();
-            } else {
-                opened = true;
             }
         });
 
-        ws.on("disconnect", () => {
+        ws.addEventListener("error", () => {
+            if(ws) ws.close();
+            manually_fetch_queue();
+        });
+
+        ws.addEventListener("close", () => {
+            setTimeout(() => {
+                ws = null;
+                createSocket(true);
+            }, 300);
             if(opened) {
                 updateStatus(null);
                 showModal("Reconnecting...");
             }
         });
 
-        ws.on("status", (data) => {
-            updateStatus(data);
-        });
-
-        ws.on("progress", (data) => {
-            updateNodeProgress(data);
-        });
-
-        ws.on("executing", (data) => {
-            setRunningNode(data.node);
-        });
-
-        ws.on("executed", (data) => {
-            nodeOutputs[data.node] = data.output;
+        ws.addEventListener("message", (event) => {
+            try {
+                const msg = JSON.parse(event.data);
+                console.log(msg.type, msg.data);
+                switch(msg.type) {
+                    case "status":
+                        if(msg.data.sid) {
+                            clientId = msg.data.sid;
+                        }
+                        updateStatus(msg.data.status);
+                        break;
+                    case "progress":
+                        updateNodeProgress(msg.data)
+                        break;
+                    case "executing":
+                        setRunningNode(msg.data.node);
+                        break;
+                    case "executed":
+                         nodeOutputs[msg.data.node] = msg.data.output;
+                        break;
+                    default:
+                        throw new Error("Unknown message type")
+                }
+            } catch (error) {
+                console.warn("Unhandled message:", event.data)
+            }
         });
     }
     createSocket();
 })();
 
+
 function clearGraph() {
     graph.clear();
 }