[Experimental] reuse established websockets for multiple image gens

for #380

Author: mcmonkey4eva

src/WebAPI/T2IAPI.cs

@@ -4,7 +4,6 @@
 using SixLabors.ImageSharp;
 using SixLabors.ImageSharp.Processing;
 using SwarmUI.Accounts;
-using SwarmUI.Backends;
 using SwarmUI.Core;
 using SwarmUI.Text2Image;
 using SwarmUI.Utils;
@@ -80,7 +79,68 @@ public static async Task<JObject> GenerateText2ImageWS(WebSocket socket, Session
         [API.APIParameter("The number of images to generate.")] int images,
         [API.APIParameter("Raw mapping of input should contain general T2I parameters (see listing on Generate tab of main interface) to values, eg `{ \"prompt\": \"a photo of a cat\", \"model\": \"OfficialStableDiffusion/sd_xl_base_1.0\", \"steps\": 20, ... }`. Note that this is the root raw map, ie all params go on the same level as `images`, `session_id`, etc.")] JObject rawInput)
     {
-        await API.RunWebsocketHandlerCallWS(GenT2I_Internal, session, (images, rawInput), socket);
+        using CancellationTokenSource cancelTok = new();
+        bool retain = false, ended = false;
+        using CancellationTokenSource linked = CancellationTokenSource.CreateLinkedTokenSource(Program.GlobalProgramCancel, cancelTok.Token);
+        SharedGenT2IData data = new();
+        ConcurrentDictionary<Task, Task> tasks = [];
+        static int guessBatchSize(JObject input)
+        {
+            if (input.TryGetValue("batchsize", out JToken batch))
+            {
+                return batch.Value<int>();
+            }
+            return 1;
+        }
+        _ = Utilities.RunCheckedTask(async () =>
+        {
+            try
+            {
+                int batchOffset = images * guessBatchSize(rawInput);
+                while (!cancelTok.IsCancellationRequested)
+                {
+                    byte[] rec = await socket.ReceiveData(1024 * 1024 * 256, linked.Token);
+                    Volatile.Write(ref retain, true);
+                    if (socket.State != WebSocketState.Open || cancelTok.IsCancellationRequested || Volatile.Read(ref ended))
+                    {
+                        return;
+                    }
+                    JObject newInput = StringConversionHelper.UTF8Encoding.GetString(rec).ParseToJson();
+                    int newImages = newInput.Value<int>("images");
+                    Task handleMore = API.RunWebsocketHandlerCallWS(GenT2I_Internal, session, (newImages, newInput, data, batchOffset), socket);
+                    tasks.TryAdd(handleMore, handleMore);
+                    Volatile.Write(ref retain, false);
+                    batchOffset += newImages * guessBatchSize(newInput);
+                }
+            }
+            catch (TaskCanceledException)
+            {
+                return;
+            }
+            finally
+            {
+                Volatile.Write(ref retain, false);
+            }
+        });
+        Task handle = API.RunWebsocketHandlerCallWS(GenT2I_Internal, session, (images, rawInput, data, 0), socket);
+        tasks.TryAdd(handle, handle);
+        while (Volatile.Read(ref retain) || tasks.Any())
+        {
+            await Task.WhenAny(tasks.Keys.ToList());
+            foreach (Task t in tasks.Keys.Where(t => t.IsCompleted).ToList())
+            {
+                tasks.TryRemove(t, out _);
+            }
+            if (tasks.IsEmpty())
+            {
+                await socket.SendJson(new JObject() { ["socket_intention"] = "close" }, API.WebsocketTimeout);
+                await Task.Delay(TimeSpan.FromSeconds(2)); // Give 2 seconds to allow a new gen request before actually closing
+                if (tasks.IsEmpty())
+                {
+                    Volatile.Write(ref ended, true);
+                }
+            }
+        }
         await socket.SendJson(BasicAPIFeatures.GetCurrentStatusRaw(session), API.WebsocketTimeout);
         return null;
     }
@@ -100,7 +160,7 @@ public static async Task<JObject> GenerateText2Image(Session session,
         [API.APIParameter("The number of images to generate.")] int images,
         [API.APIParameter("Raw mapping of input should contain general T2I parameters (see listing on Generate tab of main interface) to values, eg `{ \"prompt\": \"a photo of a cat\", \"model\": \"OfficialStableDiffusion/sd_xl_base_1.0\", \"steps\": 20, ... }`. Note that this is the root raw map, ie all params go on the same level as `images`, `session_id`, etc.")] JObject rawInput)
     {
-        List<JObject> outputs = await API.RunWebsocketHandlerCallDirect(GenT2I_Internal, session, (images, rawInput));
+        List<JObject> outputs = await API.RunWebsocketHandlerCallDirect(GenT2I_Internal, session, (images, rawInput, new SharedGenT2IData(), 0));
         Dictionary<int, string> imageOutputs = [];
         int[] discards = null;
         foreach (JObject obj in outputs)
@@ -169,10 +229,15 @@ public static T2IParamInput RequestToParams(Session session, JObject rawInput)
         return user_input;
     }
 
+    public class SharedGenT2IData
+    {
+        public int NumExtra, NumNonReal;
+    }
+
     /// <summary>Internal route for generating images.</summary>
-    public static async Task GenT2I_Internal(Session session, (int, JObject) input, Action<JObject> output, bool isWS)
+    public static async Task GenT2I_Internal(Session session, (int, JObject, SharedGenT2IData, int) input, Action<JObject> output, bool isWS)
     {
-        (int images, JObject rawInput) = input;
+        (int images, JObject rawInput, SharedGenT2IData data, int batchOffset) = input;
         using Session.GenClaim claim = session.Claim(gens: images);
         void setError(string message)
         {
@@ -214,7 +279,6 @@ void removeDoneTasks()
         }
         int max_degrees = session.User.Restrictions.CalcMaxT2ISimultaneous;
         List<int> discard = [];
-        int numExtra = 0, numNonReal = 0;
         int batchSizeExpected = user_input.Get(T2IParamTypes.BatchSize, 1);
         void saveImage(T2IEngine.ImageOutput image, int actualIndex, T2IParamInput thisParams, string metadata)
         {
@@ -255,7 +319,7 @@ void saveImage(T2IEngine.ImageOutput image, int actualIndex, T2IParamInput thisP
             {
                 break;
             }
-            int imageIndex = i * batchSizeExpected;
+            int imageIndex = i * batchSizeExpected + batchOffset;
             T2IParamInput thisParams = user_input.Clone();
             if (!thisParams.Get(T2IParamTypes.NoSeedIncrement, false))
             {
@@ -278,12 +342,12 @@ void saveImage(T2IEngine.ImageOutput image, int actualIndex, T2IParamInput thisP
                         numCalls++;
                         if (numCalls > batchSizeExpected)
                         {
-                            actualIndex = images * batchSizeExpected + Interlocked.Increment(ref numExtra);
+                            actualIndex = images * batchSizeExpected + Interlocked.Increment(ref data.NumExtra);
                         }
                     }
                     else
                     {
-                        actualIndex = -10 - Interlocked.Increment(ref numNonReal);
+                        actualIndex = -10 - Interlocked.Increment(ref data.NumNonReal);
                     }
                     saveImage(image, actualIndex, thisParams, metadata);
                 })));

src/wwwroot/js/genpage/generatehandler.js

@@ -6,6 +6,7 @@ class GenerateHandler {
         this.totalGenRunTime = 0;
         this.validateModel = true;
         this.interrupted = -1;
+        this.socket = null;
         this.imageContainerDivId = 'current_image';
         this.imageId = 'current_image_img';
         this.progressBarHtml = `<div class="image-preview-progress-inner"><div class="image-preview-progress-overall"></div><div class="image-preview-progress-current"></div></div>`;
@@ -95,7 +96,20 @@ class GenerateHandler {
             let discardable = {};
             let timeLastGenHit = Date.now();
             let actualInput = this.getGenInput(input_overrides, input_preoverrides);
-            makeWSRequestT2I('GenerateText2ImageWS', actualInput, data => {
+            let socket = null;
+            let handleData = data => {
+                if ('socket_intention' in data && data.socket_intention == 'close') {
+                    if (this.socket == socket) {
+                        this.socket = null;
+                    }
+                    if (Object.keys(discardable).length > 0) {
+                        // clear any lingering previews
+                        for (let img of Object.values(images)) {
+                            img.div.remove();
+                        }
+                    }
+                    return;
+                }
                 if (isPreview) {
                     if (data.image) {
                         this.setCurrentImage(data.image, data.metadata, `${batch_id}_${data.batch_index}`, false, true);
@@ -197,20 +211,22 @@ class GenerateHandler {
                             this.setCurrentImage(imgs[0].image, imgs[0].metadata);
                         }
                     }
-                    if (Object.keys(discardable).length > 0) {
-                        // clear any lingering previews
-                        for (let img of Object.values(images)) {
-                            img.div.remove();
-                        }
-                    }
                 }
-            }, e => {
+            };
+            let handleError = e => {
                 console.log(`Error in GenerateText2ImageWS: ${e}, ${this.interrupted}, ${batch_id}`);
                 if (this.interrupted >= batch_id) {
                     return;
                 }
                 this.hadError(e);
-            });
+            };
+            if (this.socket && this.socket.readyState == WebSocket.OPEN) {
+                this.socket.send(JSON.stringify(actualInput));
+            }
+            else {
+                socket = makeWSRequestT2I('GenerateText2ImageWS', actualInput, handleData, handleError);
+                this.socket = socket;
+            }
         };
         if (this.validateModel) {
             if (getRequiredElementById('current_model').value == '') {

src/wwwroot/js/genpage/main.js

@@ -895,7 +895,7 @@ function updateGenCount() {
 }
 
 function makeWSRequestT2I(url, in_data, callback, errorHandle = null) {
-    makeWSRequest(url, in_data, data => {
+    return makeWSRequest(url, in_data, data => {
         if (data.status) {
             updateCurrentStatusDirect(data.status);
         }

src/wwwroot/js/site.js

@@ -129,6 +129,7 @@ function makeWSRequest(url, in_data, callback, depth = 0, errorHandle = null, on
         callback(data);
     }
     socket.onerror = errorHandle ? () => errorHandle(genericServerErrorMsg.get()) : genericServerError;
+    return socket;
 }
 
 let failedCrash = translatable(`Failed to send request to server. Did the server crash?`);