From 5d14e9b9592250cd099f57ef57b620f93e75c2b7 Mon Sep 17 00:00:00 2001
From: pythongosssss <125205205+pythongosssss@users.noreply.github.com>
Date: Sun, 12 Feb 2023 15:53:48 +0000
Subject: [PATCH 01/11] Changed HTTP Server + Added WebSockets

Moved the existing API endpoints to use aoihttp and added websocket notifications
---
 main.py | 218 +++++++++++++++++++++++++++++++++++---------------------
 1 file changed, 137 insertions(+), 81 deletions(-)

diff --git a/main.py b/main.py
index 666193b6..ba204059 100644
--- a/main.py
+++ b/main.py
@@ -5,6 +5,7 @@ import json
 import threading
 import heapq
 import traceback
+import asyncio
 
 if __name__ == "__main__":
     if '--help' in sys.argv:
@@ -25,7 +26,6 @@ if '--dont-upcast-attention' in sys.argv:
     os.environ['ATTN_PRECISION'] = "fp16"
 
 import torch
-
 import nodes
 
 def get_input_data(inputs, class_def, outputs={}, prompt={}, extra_data={}):
@@ -286,16 +286,19 @@ def prompt_worker(q):
         q.task_done(item_id)
 
 class PromptQueue:
-    def __init__(self):
+    def __init__(self, socket_handler):
+        self.socket_handler = socket_handler
         self.mutex = threading.RLock()
         self.not_empty = threading.Condition(self.mutex)
         self.task_counter = 0
         self.queue = []
         self.currently_running = {}
+        socket_handler.prompt_queue = self
 
     def put(self, item):
         with self.mutex:
             heapq.heappush(self.queue, item)
+            self.socket_handler.queue_updated(self)
             self.not_empty.notify()
 
     def get(self):
@@ -306,11 +309,13 @@ class PromptQueue:
             i = self.task_counter
             self.currently_running[i] = copy.deepcopy(item)
             self.task_counter += 1
+            self.socket_handler.queue_updated(self)
             return (item, i)
 
     def task_done(self, item_id):
         with self.mutex:
             self.currently_running.pop(item_id)
+            self.socket_handler.queue_updated(self)
 
     def get_current_queue(self):
         with self.mutex:
@@ -326,6 +331,7 @@ class PromptQueue:
     def wipe_queue(self):
         with self.mutex:
             self.queue = []
+            self.socket_handler.queue_updated(self)
 
     def delete_queue_item(self, function):
         with self.mutex:
@@ -336,35 +342,82 @@ class PromptQueue:
                     else:
                         self.queue.pop(x)
                         heapq.heapify(self.queue)
+                    self.socket_handler.queue_updated(self)
                     return True
         return False
 
-from http.server import BaseHTTPRequestHandler, HTTPServer
-
-class PromptServer(BaseHTTPRequestHandler):
-    def _set_headers(self, code=200, ct='text/html'):
-        self.send_response(code)
-        self.send_header('Content-type', ct)
-        self.end_headers()
-    def log_message(self, format, *args):
-        pass
-    def do_GET(self):
-        if self.path == "/prompt":
-            self._set_headers(ct='application/json')
-            prompt_info = {}
-            exec_info = {}
-            exec_info['queue_remaining'] = self.server.prompt_queue.get_tasks_remaining()
-            prompt_info['exec_info'] = exec_info
-            self.wfile.write(json.dumps(prompt_info).encode('utf-8'))
-        elif self.path == "/queue":
-            self._set_headers(ct='application/json')
-            queue_info = {}
-            current_queue = self.server.prompt_queue.get_current_queue()
-            queue_info['queue_running'] = current_queue[0]
-            queue_info['queue_pending'] = current_queue[1]
-            self.wfile.write(json.dumps(queue_info).encode('utf-8'))
-        elif self.path == "/object_info":
-            self._set_headers(ct='application/json')
+import aiohttp
+from aiohttp import web
+
+def get_queue_info(prompt_queue):
+    prompt_info = {}
+    exec_info = {}
+    exec_info['queue_remaining'] = prompt_queue.get_tasks_remaining()
+    prompt_info['exec_info'] = exec_info
+    return prompt_info
+
+class SocketHandler():
+    def __init__(self, loop):
+        self.connected = set()
+        self.messages = asyncio.Queue()
+        self.loop = loop
+
+    async def publish_loop(self):
+        while True:
+            msg =  await self.messages.get()
+            await self.send(msg)
+
+    def queue_updated(self, queue):
+        # This is called by the queue processing thread so we need to make it thread safe
+        loop.call_soon_threadsafe(self.messages.put_nowait, { 'type': 'status', 'status': get_queue_info(queue) })
+
+    async def send(self, message, socket = None):
+        if isinstance(message, str) == False:
+            message = json.dumps(message)
+        
+        if socket is None:
+            for ws in self.connected:
+                await ws.send_str(message)
+        else:
+            await socket.send_str(message)
+    
+    async def process(self, request):
+        ws = web.WebSocketResponse()
+        await ws.prepare(request)
+        self.connected.add(ws)
+        try:
+            # Send initial state to the new client
+            await self.send({ 'type': 'status', 'status': get_queue_info(self.prompt_queue) }, ws)
+            async for msg in ws:
+                if msg.type == aiohttp.WSMsgType.ERROR:
+                    print('ws connection closed with exception %s' % ws.exception())
+        finally:
+            self.connected.remove(ws)
+
+        return ws
+        
+class PromptServer():
+    def __init__(self, prompt_queue, socket_handler):
+        self.prompt_queue = prompt_queue
+        self.socket_handler = socket_handler
+        self.number = 0
+        self.app = web.Application()
+        routes = web.RouteTableDef()
+
+        @routes.get('/ws')
+        async def websocket_handler(request):
+            return await self.socket_handler.process(request)
+        
+        @routes.get("/")
+        async def get_root(request):
+            return aiohttp.web.HTTPFound('/index.html')
+        
+        @routes.get("/prompt")
+        async def get_prompt(request):
+            return web.json_response(get_queue_info(self.prompt_queue))
+        
+        @routes.get("/object_info")
+        async def get_object_info(request):
             out = {}
             for x in nodes.NODE_CLASS_MAPPINGS:
                 obj_class = nodes.NODE_CLASS_MAPPINGS[x]
@@ -377,40 +430,32 @@ class PromptServer(BaseHTTPRequestHandler):
                 if hasattr(obj_class, 'CATEGORY'):
                     info['category'] = obj_class.CATEGORY
                 out[x] = info
-            self.wfile.write(json.dumps(out).encode('utf-8'))
-        elif self.path[1:] in os.listdir(self.server.server_dir):
-            if self.path[1:].endswith('.css'):
-                self._set_headers(ct='text/css')
-            elif self.path[1:].endswith('.js'):
-                self._set_headers(ct='text/javascript')
-            else:
-                self._set_headers()
-            with open(os.path.join(self.server.server_dir, self.path[1:]), "rb") as f:
-                self.wfile.write(f.read())
-        else:
-            self._set_headers()
-            with open(os.path.join(self.server.server_dir, "index.html"), "rb") as f:
-                self.wfile.write(f.read())
-
-    def do_HEAD(self):
-        self._set_headers()
-
-    def do_POST(self):
-        resp_code = 200
-        out_string = ""
-        if self.path == "/prompt":
+            return web.json_response(out)
+        
+        @routes.get("/queue")
+        async def get_queue(request):
+            queue_info = {}
+            current_queue = self.prompt_queue.get_current_queue()
+            queue_info['queue_running'] = current_queue[0]
+            queue_info['queue_pending'] = current_queue[1]
+            return web.json_response(queue_info)
+        
+        @routes.post("/prompt")
+        async def post_prompt(request):
             print("got prompt")
-            data_string = self.rfile.read(int(self.headers['Content-Length']))
-            json_data = json.loads(data_string)
+            resp_code = 200
+            out_string = ""
+            json_data =  await request.json()
+
             if "number" in json_data:
                 number = float(json_data['number'])
             else:
-                number = self.server.number
+                number = self.number
                 if "front" in json_data:
                     if json_data['front']:
                         number = -number
 
-                self.server.number += 1
+                self.number += 1
             if "prompt" in json_data:
                 prompt = json_data["prompt"]
                 valid = validate_prompt(prompt)
@@ -418,46 +463,54 @@ class PromptServer(BaseHTTPRequestHandler):
                 if "extra_data" in json_data:
                     extra_data = json_data["extra_data"]
                 if valid[0]:
-                    self.server.prompt_queue.put((number, id(prompt), prompt, extra_data))
+                    self.prompt_queue.put((number, id(prompt), prompt, extra_data))
                 else:
                     resp_code = 400
                     out_string = valid[1]
                     print("invalid prompt:", valid[1])
-        elif self.path == "/queue":
-            data_string = self.rfile.read(int(self.headers['Content-Length']))
-            json_data = json.loads(data_string)
+
+            return web.Response(body=out_string, status=resp_code)
+        
+        @routes.post("/queue")
+        async def post_queue(request):
+            json_data =  await request.json()
             if "clear" in json_data:
                 if json_data["clear"]:
-                    self.server.prompt_queue.wipe_queue()
+                    self.prompt_queue.wipe_queue()
             if "delete" in json_data:
                 to_delete = json_data['delete']
                 for id_to_delete in to_delete:
                     delete_func = lambda a: a[1] == int(id_to_delete)
-                    self.server.prompt_queue.delete_queue_item(delete_func)
-
-        self._set_headers(code=resp_code)
-        self.end_headers()
-        self.wfile.write(out_string.encode('utf8'))
-        return
-
-
-def run(prompt_queue, address='', port=8188):
-    server_address = (address, port)
-    httpd = HTTPServer(server_address, PromptServer)
-    httpd.server_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "webshit")
-    httpd.prompt_queue = prompt_queue
-    httpd.number = 0
-    if server_address[0] == '':
-        addr = '0.0.0.0'
-    else:
-        addr = server_address[0]
+                    self.prompt_queue.delete_queue_item(delete_func)
+                    
+            return web.Response(status=200)
+
+        self.app.add_routes(routes)
+        self.app.add_routes([
+            web.static('/', os.path.join(os.path.dirname(os.path.realpath(__file__)), "webshit")),
+        ])
+
+async def start_server(server, address, port):
+    runner = web.AppRunner(server.app)
+    await runner.setup()
+    site = web.TCPSite(runner, address, port)
+    await site.start()
+    
+    if address == '':
+        address = '0.0.0.0'
     print("Starting server\n")
-    print("To see the GUI go to: http://{}:{}".format(addr, server_address[1]))
-    httpd.serve_forever()
+    print("To see the GUI go to: http://{}:{}".format(address, port))   
 
+async def run(prompt_queue, socket_handler, address='', port=8188):
+    server = PromptServer(prompt_queue, socket_handler)
+    await asyncio.gather(start_server(server, address, port), socket_handler.publish_loop())
 
 if __name__ == "__main__":
-    q = PromptQueue()
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+
+    socket_handler = SocketHandler(loop)
+    q = PromptQueue(socket_handler)
     threading.Thread(target=prompt_worker, daemon=True, args=(q,)).start()
     if '--listen' in sys.argv:
         address = '0.0.0.0'
@@ -471,6 +524,9 @@ if __name__ == "__main__":
     except:
         pass
 
-    run(q, address=address, port=port)
+    try:
+        loop.run_until_complete(run(q, socket_handler, address=address, port=port))
+    except KeyboardInterrupt:
+        pass
 
 

From 9d1edfde1fc7541bff596c78c06d1c014d1c5b75 Mon Sep 17 00:00:00 2001
From: pythongosssss <125205205+pythongosssss@users.noreply.github.com>
Date: Sun, 12 Feb 2023 15:54:22 +0000
Subject: [PATCH 02/11] Changed polling to use websockets

---
 webshit/index.html | 58 +++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 50 insertions(+), 8 deletions(-)

diff --git a/webshit/index.html b/webshit/index.html
index 683d9648..6605148e 100644
--- a/webshit/index.html
+++ b/webshit/index.html
@@ -54,6 +54,7 @@
 
 <canvas id='mycanvas' width='1000' height='1000' style='width: 100%; height: 100%;'></canvas>
 <script>
+
 var graph = new LGraph();
 
 var canvas = new LGraphCanvas("#mycanvas", graph);
@@ -72,7 +73,6 @@ resizeCanvas();
 // call the function when the window is resized
 window.addEventListener("resize", resizeCanvas);
 
-
 var default_graph = {"last_node_id":9,"last_link_id":9,"nodes":[{"id":7,"type":"CLIPTextEncode","pos":[413,389],"size":{"0":425.27801513671875,"1":180.6060791015625},"flags":{},"order":3,"mode":0,"inputs":[{"name":"clip","type":"CLIP","link":5}],"outputs":[{"name":"CONDITIONING","type":"CONDITIONING","links":[6],"slot_index":0}],"properties":{},"widgets_values":["bad hands"]},{"id":6,"type":"CLIPTextEncode","pos":[415,186],"size":{"0":422.84503173828125,"1":164.31304931640625},"flags":{},"order":2,"mode":0,"inputs":[{"name":"clip","type":"CLIP","link":3}],"outputs":[{"name":"CONDITIONING","type":"CONDITIONING","links":[4],"slot_index":0}],"properties":{},"widgets_values":["masterpiece best quality girl"]},{"id":5,"type":"EmptyLatentImage","pos":[473,609],"size":{"0":315,"1":106},"flags":{},"order":1,"mode":0,"outputs":[{"name":"LATENT","type":"LATENT","links":[2],"slot_index":0}],"properties":{},"widgets_values":[512,512,1]},{"id":3,"type":"KSampler","pos":[863,186],"size":{"0":315,"1":262},"flags":{},"order":4,"mode":0,"inputs":[{"name":"model","type":"MODEL","link":1},{"name":"positive","type":"CONDITIONING","link":4},{"name":"negative","type":"CONDITIONING","link":6},{"name":"latent_image","type":"LATENT","link":2}],"outputs":[{"name":"LATENT","type":"LATENT","links":[7],"slot_index":0}],"properties":{},"widgets_values":[8566257,true,20,8,"sample_euler","normal",1]},{"id":8,"type":"VAEDecode","pos":[1209,188],"size":{"0":210,"1":46},"flags":{},"order":5,"mode":0,"inputs":[{"name":"samples","type":"LATENT","link":7},{"name":"vae","type":"VAE","link":8}],"outputs":[{"name":"IMAGE","type":"IMAGE","links":[9],"slot_index":0}],"properties":{}},{"id":9,"type":"SaveImage","pos":[1451,189],"size":{"0":210,"1":26},"flags":{},"order":6,"mode":0,"inputs":[{"name":"images","type":"IMAGE","link":9}],"properties":{}},{"id":4,"type":"CheckpointLoader","pos":[26,474],"size":{"0":315,"1":122},"flags":{},"order":0,"mode":0,"outputs":[{"name":"MODEL","type":"MODEL","links":[1],"slot_index":0},{"name":"CLIP","type":"CLIP","links":[3,5],"slot_index":1},{"name":"VAE","type":"VAE","links":[8],"slot_index":2}],"properties":{},"widgets_values":["v1-inference.yaml","v1-5-pruned-emaonly.ckpt"]}],"links":[[1,4,0,3,0,"MODEL"],[2,5,0,3,3,"LATENT"],[3,4,1,6,0,"CLIP"],[4,6,0,3,1,"CONDITIONING"],[5,4,1,7,0,"CLIP"],[6,7,0,3,2,"CONDITIONING"],[7,3,0,8,0,"LATENT"],[8,4,2,8,1,"VAE"],[9,8,0,9,0,"IMAGE"]],"groups":[],"config":{},"extra":{},"version":0.4}
 
 function loadGraphData(graph, graph_data)
@@ -466,14 +466,56 @@ document.addEventListener('drop', (event) => {
     prompt_file_load(file);
 });
 
+(() => {
+    function updateStatus(data) {
+        document.getElementById("queuesize").innerHTML = "Queue size: " + (data ? data.exec_info.queue_remaining : "ERR");
+    }
+
+    let ws;
+    function createSocket(isReconnect) {
+        if(ws) return;
 
-setInterval(function(){
-  fetch('/prompt')
-    .then(response => response.json())
-    .then(data => {
-        document.getElementById("queuesize").innerHTML = "Queue size: " + data.exec_info.queue_remaining + "";
-    }).catch((response) => {document.getElementById("queuesize").innerHTML = "Queue size: ERR"});
-}, 500);
+        let opened = false;
+        ws = new WebSocket(`ws://${location.host}/ws`);   
+
+        ws.addEventListener("open", () => {
+            opened = true;
+            if(isReconnect) {
+                closeModal();
+            }
+        });
+
+        ws.addEventListener("error", () => {
+            if(ws) ws.close();
+        });
+
+        ws.addEventListener("close", () => {
+            setTimeout(() => {
+                ws = null;
+                createSocket(true);
+            }, 1000);
+            if(opened) {
+                showModal("Reconnecting...");
+            }
+        });
+
+        ws.addEventListener("message", (event) => {
+            try {
+                const msg = JSON.parse(event.data);
+                switch(msg.type) {
+                    case "status":
+                        updateStatus(msg.status)
+                        break;
+                    default:
+                        throw new Error("Unknown message type")
+                }
+            } catch (error) {
+                console.warn("Unhandled message:", event.data)
+            }
+        });
+    }
+    createSocket();
+})();
 
 function clearGraph() {
     graph.clear();

From 639d21e5f78bf2e8e18515c0f49fba507cf331ac Mon Sep 17 00:00:00 2001
From: pythongosssss <125205205+pythongosssss@users.noreply.github.com>
Date: Sun, 12 Feb 2023 18:31:26 +0000
Subject: [PATCH 03/11] Added aiohttp requirement

---
 requirements.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/requirements.txt b/requirements.txt
index cc59cf1a..e4be9ebc 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -7,6 +7,6 @@ open-clip-torch
 transformers
 safetensors
 pytorch_lightning
-
+aiohttp
 accelerate
 

From f616dcc81a3e63536bffc6d250118baa3d230e5f Mon Sep 17 00:00:00 2001
From: pythongosssss <125205205+pythongosssss@users.noreply.github.com>
Date: Sun, 12 Feb 2023 18:32:20 +0000
Subject: [PATCH 04/11] Added handling of missing aiohttp module

---
 main.py | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/main.py b/main.py
index ba204059..0c1b37c7 100644
--- a/main.py
+++ b/main.py
@@ -7,6 +7,16 @@ import heapq
 import traceback
 import asyncio
 
+try:
+    import aiohttp
+    from aiohttp import web
+except ImportError:
+    print("Module 'aiohttp' not installed. Please install it via:")
+    print("pip install aiohttp")
+    print("or")
+    print("pip install -r requirements.txt")
+    sys.exit()
+
 if __name__ == "__main__":
     if '--help' in sys.argv:
         print("Valid Command line Arguments:")
@@ -346,9 +356,6 @@ class PromptQueue:
                     return True
         return False
 
-import aiohttp
-from aiohttp import web
-
 def get_queue_info(prompt_queue):
     prompt_info = {}
     exec_info = {}

From e3451cea4f2f148c48a4d3993f9e83161ac65667 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Mon, 13 Feb 2023 12:29:21 -0500
Subject: [PATCH 05/11] uni_pc now works with KSamplerAdvanced
 return_with_leftover_noise.

---
 comfy/extra_samplers/uni_pc.py | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/comfy/extra_samplers/uni_pc.py b/comfy/extra_samplers/uni_pc.py
index 42ee1c93..ae3a544e 100644
--- a/comfy/extra_samplers/uni_pc.py
+++ b/comfy/extra_samplers/uni_pc.py
@@ -821,7 +821,13 @@ def expand_dims(v, dims):
 
 
 def sample_unipc(model, noise, image, sigmas, sampling_function, extra_args=None, callback=None, disable=None):
-        timesteps = torch.nn.functional.interpolate(sigmas[None,None,:-1], size=(len(sigmas),), mode='linear')[0][0]
+        to_zero = False
+        if sigmas[-1] == 0:
+            timesteps = torch.nn.functional.interpolate(sigmas[None,None,:-1], size=(len(sigmas),), mode='linear')[0][0]
+            to_zero = True
+        else:
+            timesteps = sigmas.clone()
+
         for s in range(timesteps.shape[0]):
             timesteps[s] = (model.sigma_to_t(timesteps[s]) / 1000) + (1 / len(model.sigmas))
 
@@ -832,7 +838,7 @@ def sample_unipc(model, noise, image, sigmas, sampling_function, extra_args=None
         else:
             img = noise
 
-        if sigmas[-1] == 0:
+        if to_zero:
             timesteps[-1] = (1 / len(model.sigmas))
 
         device = noise.device
@@ -853,4 +859,6 @@ def sample_unipc(model, noise, image, sigmas, sampling_function, extra_args=None
 
         uni_pc = UniPC(model_fn, ns, predict_x0=True, thresholding=False)
         x = uni_pc.sample(img, timesteps=timesteps, skip_type="time_uniform", method="multistep", order=3, lower_order_final=True)
+        if not to_zero:
+            x /= ns.marginal_alpha(timesteps[-1])
         return x

From 540bbbdb825fe1b02557da2658b7cbe9e643c707 Mon Sep 17 00:00:00 2001
From: pythongosssss <125205205+pythongosssss@users.noreply.github.com>
Date: Mon, 13 Feb 2023 17:53:21 +0000
Subject: [PATCH 06/11] Handle https when creating socket

---
 webshit/index.html | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/webshit/index.html b/webshit/index.html
index d5027665..c3f56afc 100644
--- a/webshit/index.html
+++ b/webshit/index.html
@@ -490,7 +490,7 @@ document.addEventListener('drop', (event) => {
         if(ws) return;
 
         let opened = false;
-        ws = new WebSocket(`ws://${location.host}/ws`);   
+        ws = new WebSocket(`ws${window.location.protocol === "https"? "s" : ""}://${location.host}/ws`);   
 
         ws.addEventListener("open", () => {
             opened = true;

From 8e51af6786aa4132240d3b1d8b59aa1f7552e24a Mon Sep 17 00:00:00 2001
From: pythongosssss <125205205+pythongosssss@users.noreply.github.com>
Date: Mon, 13 Feb 2023 17:58:03 +0000
Subject: [PATCH 07/11] Changed to serve index.html on root

---
 main.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/main.py b/main.py
index 0c1b37c7..7de266a8 100644
--- a/main.py
+++ b/main.py
@@ -409,6 +409,7 @@ class PromptServer():
         self.socket_handler = socket_handler
         self.number = 0
         self.app = web.Application()
+        self.web_root = os.path.join(os.path.dirname(os.path.realpath(__file__)), "webshit")
         routes = web.RouteTableDef()
 
         @routes.get('/ws')
@@ -417,7 +418,7 @@ class PromptServer():
         
         @routes.get("/")
         async def get_root(request):
-            return aiohttp.web.HTTPFound('/index.html')
+            return web.FileResponse(os.path.join(self.web_root, "index.html"))
         
         @routes.get("/prompt")
         async def get_prompt(request):
@@ -494,7 +495,7 @@ class PromptServer():
 
         self.app.add_routes(routes)
         self.app.add_routes([
-            web.static('/', os.path.join(os.path.dirname(os.path.realpath(__file__)), "webshit")),
+            web.static('/', self.web_root),
         ])
 
 async def start_server(server, address, port):

From 68a1ac774611c4df3e45f8e99f26ed483e3a0306 Mon Sep 17 00:00:00 2001
From: pythongosssss <125205205+pythongosssss@users.noreply.github.com>
Date: Mon, 13 Feb 2023 20:12:12 +0000
Subject: [PATCH 08/11] Remove catch that was causing linux ctrl+c issues

This was for hiding an error on Windows, but it isn't required
---
 main.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/main.py b/main.py
index 7de266a8..c5f04aa9 100644
--- a/main.py
+++ b/main.py
@@ -532,9 +532,5 @@ if __name__ == "__main__":
     except:
         pass
 
-    try:
-        loop.run_until_complete(run(q, socket_handler, address=address, port=port))
-    except KeyboardInterrupt:
-        pass
-
+    loop.run_until_complete(run(q, socket_handler, address=address, port=port))
 

From 93d84e3b9c21126b6d6a5911a43d924348f9f2df Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Mon, 13 Feb 2023 15:52:34 -0500
Subject: [PATCH 09/11] Lower timeout and fix for colab.

---
 webshit/index.html | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/webshit/index.html b/webshit/index.html
index c3f56afc..66d09660 100644
--- a/webshit/index.html
+++ b/webshit/index.html
@@ -480,11 +480,21 @@ document.addEventListener('drop', (event) => {
     prompt_file_load(file);
 });
 
+
 (() => {
     function updateStatus(data) {
         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.exec_info.queue_remaining);
+            }).catch((response) => {updateStatus(null)});
+    }
+
     let ws;
     function createSocket(isReconnect) {
         if(ws) return;
@@ -501,14 +511,16 @@ document.addEventListener('drop', (event) => {
 
         ws.addEventListener("error", () => {
             if(ws) ws.close();
+            manually_fetch_queue();
         });
 
         ws.addEventListener("close", () => {
             setTimeout(() => {
                 ws = null;
                 createSocket(true);
-            }, 1000);
+            }, 300);
             if(opened) {
+                updateStatus(null);
                 showModal("Reconnecting...");
             }
         });

From 73251f92c830fe7f85ef284d5e1d6db73be0afca Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Mon, 13 Feb 2023 16:04:26 -0500
Subject: [PATCH 10/11] Now it works in colab properly again.

---
 webshit/index.html | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/webshit/index.html b/webshit/index.html
index 66d09660..d9b11cfe 100644
--- a/webshit/index.html
+++ b/webshit/index.html
@@ -491,7 +491,7 @@ document.addEventListener('drop', (event) => {
         fetch('/prompt')
             .then(response => response.json())
             .then(data => {
-                updateStatus(data.exec_info.queue_remaining);
+                updateStatus(data);
             }).catch((response) => {updateStatus(null)});
     }
 

From c1d5810020ba63dfa2778c0b473aa8570cea6d03 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Mon, 13 Feb 2023 16:07:34 -0500
Subject: [PATCH 11/11] Another small fix.

---
 webshit/index.html | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/webshit/index.html b/webshit/index.html
index d9b11cfe..3e2645a2 100644
--- a/webshit/index.html
+++ b/webshit/index.html
@@ -500,7 +500,7 @@ document.addEventListener('drop', (event) => {
         if(ws) return;
 
         let opened = false;
-        ws = new WebSocket(`ws${window.location.protocol === "https"? "s" : ""}://${location.host}/ws`);   
+        ws = new WebSocket(`ws${window.location.protocol === "https:"? "s" : ""}://${location.host}/ws`);
 
         ws.addEventListener("open", () => {
             opened = true;