67 lines
2.7 KiB
JavaScript
67 lines
2.7 KiB
JavaScript
import { app } from "/scripts/app.js";
|
|
import { api } from "/scripts/api.js";
|
|
|
|
app.registerExtension({
|
|
name: "ComfyUI.FlowControl.PersistentQueue",
|
|
async setup() {
|
|
const patchRgthree = async () => {
|
|
try {
|
|
const rgthree = await import("/extensions/rgthree-comfy/common/prompt_service.js");
|
|
if (rgthree && rgthree.SERVICE) {
|
|
const queue = await api.getQueue();
|
|
const allItems = [...(queue.Running || []), ...(queue.Pending || [])];
|
|
let patched = false;
|
|
for (const item of allItems) {
|
|
const promptId = item[1];
|
|
const promptData = item[2];
|
|
if (promptId && promptData) {
|
|
const promptExecution = rgthree.SERVICE.getOrMakePrompt(promptId);
|
|
if (!promptExecution.totalNodes) {
|
|
promptExecution.setPrompt({ output: promptData });
|
|
patched = true;
|
|
}
|
|
}
|
|
}
|
|
if (patched) {
|
|
rgthree.SERVICE.dispatchProgressUpdate();
|
|
console.log("[FlowControl] Synced recovered queue with rgthree-comfy.");
|
|
}
|
|
}
|
|
} catch (e) {
|
|
// rgthree-comfy not installed or path changed, ignore
|
|
}
|
|
};
|
|
|
|
const claimQueueOwnership = async () => {
|
|
if (api && api.clientId) {
|
|
try {
|
|
await api.fetchApi("/persistent_queue/claim", {
|
|
method: "POST",
|
|
headers: { "Content-Type": "application/json" },
|
|
body: JSON.stringify({ client_id: api.clientId })
|
|
});
|
|
|
|
// Wait for server to sync queue state back to clients, then patch rgthree
|
|
setTimeout(patchRgthree, 500);
|
|
} catch (e) {
|
|
console.error("[FlowControl] Error claiming queue ownership:", e);
|
|
}
|
|
}
|
|
};
|
|
|
|
// Delay execution until after app & API client ID initialization
|
|
setTimeout(claimQueueOwnership, 1000);
|
|
|
|
if (api && api.addEventListener) {
|
|
api.addEventListener("status", () => {
|
|
claimQueueOwnership();
|
|
}, { once: true });
|
|
|
|
// Also run patch anytime an execution starts just in case it was triggered externally
|
|
api.addEventListener("execution_start", () => {
|
|
setTimeout(patchRgthree, 100);
|
|
});
|
|
}
|
|
}
|
|
});
|