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] 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");
+    }
 
-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 ws;
+    function createSocket(isReconnect) {
+        if(ws) return;
+
+        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();