Fix Loopy bug with queues

AshMartian · Feb 19, 2024 · d46e29a · d46e29a
1 parent 6036865
commit d46e29a
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 5 deletions.
diff --git a/README.md b/README.md
@@ -29,7 +29,7 @@ Then restart ComfyUI
 |                                 Node                                  | Description                                                                                                                                                                                                                                                                                                                                                                               |
 | :-------------------------------------------------------------------: | :---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
 | GIR Directory Picker![GIR Directory Picker](docs/DirectoryPicker.png) | **"Select Directory"** opens a GUI file browser using `tkinker`, the selected directory is persistent on disk (server restarts are no problem!), and available via output string. Changing the Selected Directory manually is supported too.                                                                                                                                              |
-|         GIR Loopy Dir![GIR Loopy Dir](docs/DirectoryLoop.png)         | Directory in, automatic loopy-ness out. Auto-increment and auto resetting, when the `loop_index` reaches the end of filtered files in directory, goes back to zero. Filter by extension or regex, manually change `loop_index` to override!                                                                                                                                               |
+|         GIR Loopy Dir![GIR Loopy Dir](docs/DirectoryLoop.png)         | Directory in, automatic loopy-ness out. Auto-increment and auto resetting, when the `loop_index` reaches the end of filtered files in directory, goes back to zero. Filter by extension or regex, manually change `loop_index` to override! Note: When queuing, the `loop_index` will increment in real time, not pre-generated like seed inputs.                                                                                                                                              |
 |      GIR Image Nabber ![GIR Image Nabber](docs/ImageNabber.png)       | A simple path to image node, for working with results of Loopy Dir without relying on an external node pack that has it's own directory filtering logic _(For videos, use [VHS Load Video (Path)](https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite/blob/main/videohelpersuite/load_video_nodes.py#L197)_)                                                                          |
 |        GIR Happy Dance!\* ![Let's Dance!](docs/GIR_Dance.gif)         | These nodes simplify and streamline the process of directory selection and iteration, making complex tasks feel like a breeze. With the ability to effortlessly pick directories and loop through files with smart reset capabilities, GIR is over the moon, knowing users can focus on creativity and productivity, leaving the tedious parts to the automation magic of ComfyUI DirGir. |
 

diff --git a/dir_loop.py b/dir_loop.py
@@ -89,3 +89,9 @@ def iterate_directory(cls, directory, filter_type, filter_value, loop_index, pro
 @server.PromptServer.instance.routes.get("/gir-dir/loop-index")
 async def get_last_index(request):
     return web.json_response({'loop_index': loop_indexes.get(request.rel_url.query.get('id', '')) or 0})
+
+
+@server.PromptServer.instance.routes.get("/gir-dir/loop-queue")
+async def get_loop_queue(request):
+    loop_indexes[request.rel_url.query.get('id', '')] += 1
+    return web.json_response({'loop_queue': loop_indexes.get(request.rel_url.query.get('id', '')) or 0})
diff --git a/web/js/gir.js b/web/js/gir.js
@@ -18,7 +18,7 @@ function getDirectoryPath(node_id) {
         resolve(data.selected_folder);
     })
     .catch(error => {
-        console.error('Error setting directory:', error);
+        console.error('Error selecting directory:', error);
     });
   });
 }
@@ -36,7 +36,7 @@ function getCurrentDirectory(node_id) {
         resolve(data.selected_folder);
     })
     .catch(error => {
-        console.error('Error setting directory:', error);
+        console.error('Error getting directory:', error);
         reject(error);
     });
   });
@@ -73,7 +73,26 @@ function getLoopIndex(node_id) {
         resolve(data.loop_index);
     })
     .catch(error => {
-        console.error('Error setting directory:', error);
+        console.error('Error getting loop index:', error);
+        reject(error);
+    });
+  });
+}
+
+function queueLoopIndex(node_id) {
+  return new Promise((resolve, reject) => {
+  api.fetchApi('/gir-dir/queue-loop-index?id=' + node_id)
+    .then(response => {
+        if (!response.ok) {
+            throw new Error('Network response was not ok');
+        }
+        return response.json();
+    })
+    .then(data => {
+        resolve(data.loop_index);
+    })
+    .catch(error => {
+        console.error('Error triggering queue loop:', error);
         reject(error);
     });
   });
@@ -159,7 +178,21 @@ app.registerExtension({
           node.widgets.find(w => w.name === 'loop_index').value = loop_index;
         });
 
-        api.addEventListener('executed', async () => {
+        let lastQueueRemaining = 0;
+        api.addEventListener('status', async (status) => {
+          if (status.detail.exec_info.queue_remaining === 0) {
+            getLoopIndex(node.id).then(loop_index => {
+              node.widgets.find(w => w.name === 'loop_index').value = loop_index;
+            });
+          } else if (status.detail.exec_info.queue_remaining !== lastQueueRemaining) {
+            // Trigger a queue loop index
+            queueLoopIndex(node.id).then(loop_index => {
+              node.widgets.find(w => w.name === 'loop_index').value = loop_index;
+            });
+          }
+          lastQueueRemaining = status.detail.exec_info.queue_remaining;
+        });
+        api.addEventListener('executed', async (status) => {
           getLoopIndex(node.id).then(loop_index => {
             node.widgets.find(w => w.name === 'loop_index').value = loop_index;
           });