easydiffusion/ui/media/js/task-manager.js

394 lines
15 KiB
JavaScript

const htmlTaskMap = new WeakMap()
const pauseBtn = document.querySelector("#pause")
const resumeBtn = document.querySelector("#resume")
const processOrder = document.querySelector("#process_order_toggle")
let TASK_CALLBACKS = {
before_task_start: [],
after_task_start: [],
on_task_step: [],
on_render_task_success: [],
on_render_task_fail: [],
on_all_tasks_complete: [],
}
let pauseClient = false
async function onIdle() {
const serverCapacity = SD.serverCapacity
if (pauseClient === true) {
await resumeClient()
}
for (const taskEntry of getUncompletedTaskEntries()) {
if (SD.activeTasks.size >= serverCapacity) {
break
}
const task = htmlTaskMap.get(taskEntry)
if (!task) {
const taskStatusLabel = taskEntry.querySelector(".taskStatusLabel")
taskStatusLabel.style.display = "none"
continue
}
await onTaskStart(task)
}
}
function getUncompletedTaskEntries() {
const taskEntries = Array.from(document.querySelectorAll("#preview .imageTaskContainer .taskStatusLabel"))
.filter((taskLabel) => taskLabel.style.display !== "none")
.map(function(taskLabel) {
let imageTaskContainer = taskLabel.parentNode
while (!imageTaskContainer.classList.contains("imageTaskContainer") && imageTaskContainer.parentNode) {
imageTaskContainer = imageTaskContainer.parentNode
}
return imageTaskContainer
})
if (!processOrder.checked) {
taskEntries.reverse()
}
return taskEntries
}
async function onTaskStart(task) {
if (!task.isProcessing || task.batchesDone >= task.batchCount) {
return
}
if (typeof task.startTime !== "number") {
task.startTime = Date.now()
}
if (!("instances" in task)) {
task["instances"] = []
}
task["stopTask"].innerHTML = '<i class="fa-solid fa-circle-stop"></i> Stop'
task["taskStatusLabel"].innerText = "Starting"
task["taskStatusLabel"].classList.add("waitingTaskLabel")
if (task.previewTaskReq !== undefined) {
let controlImagePreview = task.taskConfig.querySelector(".controlnet-img-preview > img")
try {
let result = await SD.filter(task.previewTaskReq)
controlImagePreview.src = result.output[0]
let controlImageLargePreview = task.taskConfig.querySelector(
".controlnet-img-preview .task-fs-initimage img"
)
controlImageLargePreview.src = controlImagePreview.src
} catch (error) {
console.log("filter error", error)
}
delete task.previewTaskReq
}
let newTaskReqBody = task.reqBody
if (task.batchCount > 1) {
// Each output render batch needs it's own task reqBody instance to avoid altering the other runs after they are completed.
newTaskReqBody = Object.assign({}, task.reqBody)
if (task.batchesDone == task.batchCount - 1) {
// Last batch of the task
// If the number of parallel jobs is no factor of the total number of images, the last batch must create less than "parallel jobs count" images
// E.g. with numOutputsTotal = 6 and num_outputs = 5, the last batch shall only generate 1 image.
newTaskReqBody.num_outputs = task.numOutputsTotal - task.reqBody.num_outputs * (task.batchCount - 1)
}
}
const startSeed = task.seed || newTaskReqBody.seed
const genSeeds = Boolean(
typeof newTaskReqBody.seed !== "number" || (newTaskReqBody.seed === task.seed && task.numOutputsTotal > 1)
)
if (genSeeds) {
newTaskReqBody.seed = parseInt(startSeed) + task.batchesDone * task.reqBody.num_outputs
}
const outputContainer = document.createElement("div")
outputContainer.className = "img-batch"
task.outputContainer.insertBefore(outputContainer, task.outputContainer.firstChild)
const eventInfo = { reqBody: newTaskReqBody }
const callbacksPromises = PLUGINS["TASK_CREATE"].map((hook) => {
if (typeof hook !== "function") {
console.error("The provided TASK_CREATE hook is not a function. Hook: %o", hook)
return Promise.reject(new Error("hook is not a function."))
}
try {
return Promise.resolve(hook.call(task, eventInfo))
} catch (err) {
console.error(err)
return Promise.reject(err)
}
})
await Promise.allSettled(callbacksPromises)
let instance = eventInfo.instance
if (!instance) {
const factory = PLUGINS.OUTPUTS_FORMATS.get(eventInfo.reqBody?.output_format || newTaskReqBody.output_format)
if (factory) {
instance = await Promise.resolve(factory(eventInfo.reqBody || newTaskReqBody))
}
if (!instance) {
console.error(
`${factory ? "Factory " + String(factory) : "No factory defined"} for output format ${eventInfo.reqBody
?.output_format || newTaskReqBody.output_format}. Instance is ${instance ||
"undefined"}. Using default renderer.`
)
instance = new SD.RenderTask(eventInfo.reqBody || newTaskReqBody)
}
}
task["instances"].push(instance)
task.batchesDone++
TASK_CALLBACKS["before_task_start"].forEach((callback) => callback(task))
instance.enqueue(getTaskUpdater(task, newTaskReqBody, outputContainer)).then(
(renderResult) => {
onRenderTaskCompleted(task, newTaskReqBody, instance, outputContainer, renderResult)
},
(reason) => {
onTaskErrorHandler(task, newTaskReqBody, instance, reason)
}
)
TASK_CALLBACKS["after_task_start"].forEach((callback) => callback(task))
}
function getTaskUpdater(task, reqBody, outputContainer) {
const outputMsg = task["outputMsg"]
const progressBar = task["progressBar"]
const progressBarInner = progressBar.querySelector("div")
const batchCount = task.batchCount
let lastStatus = undefined
return async function(event) {
if (this.status !== lastStatus) {
lastStatus = this.status
switch (this.status) {
case SD.TaskStatus.pending:
task["taskStatusLabel"].innerText = "Pending"
task["taskStatusLabel"].classList.add("waitingTaskLabel")
break
case SD.TaskStatus.waiting:
task["taskStatusLabel"].innerText = "Waiting"
task["taskStatusLabel"].classList.add("waitingTaskLabel")
task["taskStatusLabel"].classList.remove("activeTaskLabel")
break
case SD.TaskStatus.processing:
case SD.TaskStatus.completed:
task["taskStatusLabel"].innerText = "Processing"
task["taskStatusLabel"].classList.add("activeTaskLabel")
task["taskStatusLabel"].classList.remove("waitingTaskLabel")
break
case SD.TaskStatus.stopped:
break
case SD.TaskStatus.failed:
if (!SD.isServerAvailable()) {
logError(
"Stable Diffusion is still starting up, please wait. If this goes on beyond a few minutes, Stable Diffusion has probably crashed. Please check the error message in the command-line window.",
event,
outputMsg
)
} else if (typeof event?.response === "object") {
let msg = "Stable Diffusion had an error reading the response:<br/><pre>"
if (this.exception) {
msg += `Error: ${this.exception.message}<br/>`
}
try {
// 'Response': body stream already read
msg += "Read: " + (await event.response.text())
} catch (e) {
msg += "Unexpected end of stream. "
}
const bufferString = event.reader.bufferedString
if (bufferString) {
msg += "Buffered data: " + bufferString
}
msg += "</pre>"
logError(msg, event, outputMsg)
}
break
}
}
if ("update" in event) {
const stepUpdate = event.update
if (!("step" in stepUpdate)) {
return
}
// task.instances can be a mix of different tasks with uneven number of steps (Render Vs Filter Tasks)
const instancesWithProgressUpdates = task.instances.filter((instance) => instance.step !== undefined)
const overallStepCount =
instancesWithProgressUpdates.reduce(
(sum, instance) =>
sum +
(instance.isPending
? Math.max(0, instance.step || stepUpdate.step) /
(instance.total_steps || stepUpdate.total_steps)
: 1),
0 // Initial value
) * stepUpdate.total_steps // Scale to current number of steps.
const totalSteps = instancesWithProgressUpdates.reduce(
(sum, instance) => sum + (instance.total_steps || stepUpdate.total_steps),
stepUpdate.total_steps * (batchCount - task.batchesDone) // Initial value at (unstarted task count * Nbr of steps)
)
const percent = Math.min(100, 100 * (overallStepCount / totalSteps)).toFixed(0)
const timeTaken = stepUpdate.step_time // sec
const stepsRemaining = Math.max(0, totalSteps - overallStepCount)
const timeRemaining = timeTaken < 0 ? "" : millisecondsToStr(stepsRemaining * timeTaken * 1000)
outputMsg.innerHTML = `Batch ${task.batchesDone} of ${batchCount}. Generating image(s): ${percent}%. Time remaining (approx): ${timeRemaining}`
outputMsg.style.display = "block"
progressBarInner.style.width = `${percent}%`
if (stepUpdate.output) {
TASK_CALLBACKS["on_task_step"].forEach((callback) =>
callback(task, reqBody, stepUpdate, outputContainer)
)
}
}
}
}
function onRenderTaskCompleted(task, reqBody, instance, outputContainer, stepUpdate) {
if (typeof stepUpdate === "object") {
if (stepUpdate.status === "succeeded") {
TASK_CALLBACKS["on_render_task_success"].forEach((callback) =>
callback(task, reqBody, stepUpdate, outputContainer)
)
} else {
task.isProcessing = false
TASK_CALLBACKS["on_render_task_fail"].forEach((callback) =>
callback(task, reqBody, stepUpdate, outputContainer)
)
}
}
if (task.isProcessing && task.batchesDone < task.batchCount) {
task["taskStatusLabel"].innerText = "Pending"
task["taskStatusLabel"].classList.add("waitingTaskLabel")
task["taskStatusLabel"].classList.remove("activeTaskLabel")
return
}
if ("instances" in task && task.instances.some((ins) => ins != instance && ins.isPending)) {
return
}
task.isProcessing = false
task["stopTask"].innerHTML = '<i class="fa-solid fa-trash-can"></i> Remove'
task["taskStatusLabel"].style.display = "none"
let time = millisecondsToStr(Date.now() - task.startTime)
if (task.batchesDone == task.batchCount) {
if (!task.outputMsg.innerText.toLowerCase().includes("error")) {
task.outputMsg.innerText = `Processed ${task.numOutputsTotal} images in ${time}`
}
task.progressBar.style.height = "0px"
task.progressBar.style.border = "0px solid var(--background-color3)"
task.progressBar.classList.remove("active")
// setStatus("request", "done", "success")
} else {
task.outputMsg.innerText += `. Task ended after ${time}`
}
// if (randomSeedField.checked) { // we already update this before the task starts
// seedField.value = task.seed
// }
if (SD.activeTasks.size > 0) {
return
}
const uncompletedTasks = getUncompletedTaskEntries()
if (uncompletedTasks && uncompletedTasks.length > 0) {
return
}
if (pauseClient) {
resumeBtn.click()
}
TASK_CALLBACKS["on_all_tasks_complete"].forEach((callback) => callback())
}
function resumeClient() {
if (pauseClient) {
document.body.classList.remove("wait-pause")
document.body.classList.add("pause")
}
return new Promise((resolve) => {
let playbuttonclick = function() {
resumeBtn.removeEventListener("click", playbuttonclick)
resolve("resolved")
}
resumeBtn.addEventListener("click", playbuttonclick)
})
}
function abortTask(task) {
if (!task.isProcessing) {
return false
}
task.isProcessing = false
task.progressBar.classList.remove("active")
task["taskStatusLabel"].style.display = "none"
task["stopTask"].innerHTML = '<i class="fa-solid fa-trash-can"></i> Remove'
if (!task.instances?.some((r) => r.isPending)) {
return
}
task.instances.forEach((instance) => {
try {
instance.abort()
} catch (e) {
console.error(e)
}
})
}
async function stopAllTasks() {
getUncompletedTaskEntries().forEach((taskEntry) => {
const taskStatusLabel = taskEntry.querySelector(".taskStatusLabel")
if (taskStatusLabel) {
taskStatusLabel.style.display = "none"
}
const task = htmlTaskMap.get(taskEntry)
if (!task) {
return
}
abortTask(task)
})
}
function onTaskErrorHandler(task, reqBody, instance, reason) {
if (!task.isProcessing) {
return
}
console.log("Render request %o, Instance: %o, Error: %s", reqBody, instance, reason)
abortTask(task)
const outputMsg = task["outputMsg"]
logError(
"Stable Diffusion had an error. Please check the logs in the command-line window. <br/><br/>" +
reason +
"<br/><pre>" +
reason.stack +
"</pre>",
task,
outputMsg
)
// setStatus("request", "error", "error")
}
pauseBtn.addEventListener("click", function() {
pauseClient = true
pauseBtn.style.display = "none"
resumeBtn.style.display = "inline"
document.body.classList.add("wait-pause")
})
resumeBtn.addEventListener("click", function() {
pauseClient = false
resumeBtn.style.display = "none"
pauseBtn.style.display = "inline"
document.body.classList.remove("pause")
document.body.classList.remove("wait-pause")
})