feat: add generation cancel job functionality and improve job handling
- Introduced `generationCancelJob` command ID in `commands/ids.ts`. - Added `GenerationCancelJobPayload` type in `commands/payloads.ts`. - Enhanced job status to include "cancelled" in `editor/state.ts`. - Updated `runGenerationJob` to accept an `AbortSignal` and handle cancellation. - Implemented cancellation logic in `runGenerate` and related functions. - Added tests for job cancellation in `operations/generation/workflow.test.ts`. - Improved layer rendering logic to prevent stack overflow in `editor/document-indexes.ts`. - Added raster size assertions in `platform/browser/rasterLimits.ts` for image processing limits. - Enhanced image file handling to check for size limits in `platform/browser/imageFiles.ts`. - Updated UI components to reflect job cancellation state in `view/GenerationJobStatus.tsx` and `view/bottom-controls/GenerateActionControls.tsx`.
This commit is contained in:
syntaxbullet
@@ -4,7 +4,7 @@ export async function handleComfyApi(request: Request) {
|
||||
try {
|
||||
const url = new URL(request.url);
|
||||
if (url.pathname === "/api/comfy/models" && request.method === "GET") return json(await listGenerationOptions());
|
||||
if (url.pathname === "/api/comfy/generate" && request.method === "POST") return json(await generate(await request.json() as ComfyGenerateRequest));
|
||||
if (url.pathname === "/api/comfy/generate" && request.method === "POST") return json(await generate(await request.json() as ComfyGenerateRequest, request.signal));
|
||||
return new Response("Not found", { status: 404 });
|
||||
} catch (error) {
|
||||
return new Response(error instanceof Error ? error.message : "ComfyUI request failed", { status: 500 });
|
||||
|
||||
@@ -111,15 +111,15 @@ async function listCheckpointModels() {
|
||||
return (await listGenerationOptions()).models;
|
||||
}
|
||||
|
||||
export async function generate(request: ComfyGenerateRequest) {
|
||||
export async function generate(request: ComfyGenerateRequest, signal?: AbortSignal) {
|
||||
if (!request.prompt?.trim()) throw new Error("Prompt is required");
|
||||
const architecture = normalizeArchitecture(request.architecture);
|
||||
if (request.mode !== "text-to-image" && architecture !== "sdxl") throw new Error(`${architectureLabel(architecture)} currently supports text-to-image only`);
|
||||
if (!request.model || request.model === "auto") request.model = await defaultModelForArchitecture(architecture);
|
||||
|
||||
const clientId = crypto.randomUUID();
|
||||
const uploaded = request.inputImage ? await uploadDataUrl(request.inputImage, `image-studio-${crypto.randomUUID()}.png`) : undefined;
|
||||
const mask = request.maskImage ? await uploadDataUrl(request.maskImage, `image-studio-mask-${crypto.randomUUID()}.png`) : undefined;
|
||||
const uploaded = request.inputImage ? await uploadDataUrl(request.inputImage, `image-studio-${crypto.randomUUID()}.png`, signal) : undefined;
|
||||
const mask = request.maskImage ? await uploadDataUrl(request.maskImage, `image-studio-mask-${crypto.randomUUID()}.png`, signal) : undefined;
|
||||
if (request.mode === "inpaint" && (!uploaded || !mask)) throw new Error("Inpaint requires normalized input and mask images");
|
||||
const prompt = buildComfyWorkflow({ ...request, architecture, inputImage: uploaded, maskImage: mask });
|
||||
|
||||
@@ -127,6 +127,7 @@ export async function generate(request: ComfyGenerateRequest) {
|
||||
method: "POST",
|
||||
headers: { "content-type": "application/json" },
|
||||
body: JSON.stringify({ client_id: clientId, prompt }),
|
||||
signal,
|
||||
});
|
||||
if (!queued.ok) throw new Error(`ComfyUI prompt failed: ${queued.status} ${await queued.text()}`);
|
||||
const queuedBody = await queued.json() as { prompt_id?: string; node_errors?: unknown };
|
||||
@@ -134,19 +135,25 @@ export async function generate(request: ComfyGenerateRequest) {
|
||||
if (nodeError) throw new Error(`ComfyUI rejected the workflow: ${nodeError}`);
|
||||
if (!queuedBody.prompt_id) throw new Error("ComfyUI did not return a prompt id");
|
||||
const prompt_id = queuedBody.prompt_id;
|
||||
const history = await waitForHistory(prompt_id);
|
||||
let history: unknown;
|
||||
try {
|
||||
history = await waitForHistory(prompt_id, signal);
|
||||
} catch (error) {
|
||||
if (signal?.aborted) await cancelComfyPrompt(prompt_id);
|
||||
throw error;
|
||||
}
|
||||
const historyError = historyErrorMessage(history);
|
||||
if (historyError) throw new Error(`ComfyUI generation failed: ${historyError}`);
|
||||
const image = selectGeneratedOutputImage(history);
|
||||
if (!image) throw new Error("ComfyUI did not return an image");
|
||||
|
||||
const imageResponse = await fetch(`${comfyBaseUrl}/view?${new URLSearchParams({ filename: image.filename, subfolder: image.subfolder ?? "", type: image.type ?? "output" })}`);
|
||||
const imageResponse = await fetch(`${comfyBaseUrl}/view?${new URLSearchParams({ filename: image.filename, subfolder: image.subfolder ?? "", type: image.type ?? "output" })}`, { signal });
|
||||
if (!imageResponse.ok) throw new Error(`ComfyUI image fetch failed: ${imageResponse.status}`);
|
||||
const bytes = Buffer.from(await imageResponse.arrayBuffer());
|
||||
return { source: `data:image/png;base64,${bytes.toString("base64")}`, mimeType: "image/png" };
|
||||
}
|
||||
|
||||
async function uploadDataUrl(dataUrl: string, filename: string) {
|
||||
async function uploadDataUrl(dataUrl: string, filename: string, signal?: AbortSignal) {
|
||||
const match = /^data:([^;]+);base64,(.+)$/.exec(dataUrl);
|
||||
if (!match) throw new Error("Expected a base64 data URL image");
|
||||
const mimeType = match[1] ?? "image/png";
|
||||
@@ -154,27 +161,53 @@ async function uploadDataUrl(dataUrl: string, filename: string) {
|
||||
const form = new FormData();
|
||||
form.append("image", new File([new Uint8Array(Buffer.from(base64, "base64"))], filename, { type: mimeType }));
|
||||
form.append("overwrite", "true");
|
||||
const response = await fetch(`${comfyBaseUrl}/upload/image`, { method: "POST", body: form });
|
||||
const response = await fetch(`${comfyBaseUrl}/upload/image`, { method: "POST", body: form, signal });
|
||||
if (!response.ok) throw new Error(`ComfyUI upload failed: ${response.status}`);
|
||||
const uploaded = await response.json() as { name: string };
|
||||
return uploaded.name;
|
||||
}
|
||||
|
||||
async function waitForHistory(promptId: string) {
|
||||
async function waitForHistory(promptId: string, signal?: AbortSignal) {
|
||||
const startedAt = Date.now();
|
||||
let attempts = 0;
|
||||
while (Date.now() - startedAt < comfyHistoryTimeoutMs) {
|
||||
const response = await fetch(`${comfyBaseUrl}/history/${promptId}`);
|
||||
if (signal?.aborted) {
|
||||
await cancelComfyPrompt(promptId);
|
||||
throw signal.reason ?? new DOMException("Generation cancelled", "AbortError");
|
||||
}
|
||||
const response = await fetch(`${comfyBaseUrl}/history/${promptId}`, { signal });
|
||||
attempts += 1;
|
||||
if (response.ok) {
|
||||
const history = await response.json() as Record<string, unknown>;
|
||||
if (history[promptId]) return history[promptId];
|
||||
}
|
||||
await Bun.sleep(comfyHistoryPollIntervalMs);
|
||||
await abortableSleep(comfyHistoryPollIntervalMs, signal);
|
||||
}
|
||||
throw new Error(`Timed out waiting for ComfyUI prompt ${promptId} after ${Math.round(comfyHistoryTimeoutMs / 1000)} seconds and ${attempts} checks`);
|
||||
}
|
||||
|
||||
async function cancelComfyPrompt(promptId: string) {
|
||||
const queueResponse = await fetch(`${comfyBaseUrl}/queue`).catch(() => undefined);
|
||||
const queue = queueResponse?.ok ? await queueResponse.json() as { queue_running?: unknown[][]; queue_pending?: unknown[][] } : undefined;
|
||||
const isRunning = queue?.queue_running?.some((entry) => entry.includes(promptId)) ?? false;
|
||||
const isPending = queue?.queue_pending?.some((entry) => entry.includes(promptId)) ?? true;
|
||||
const requests: Promise<unknown>[] = [];
|
||||
if (isPending) requests.push(fetch(`${comfyBaseUrl}/queue`, { method: "POST", headers: { "content-type": "application/json" }, body: JSON.stringify({ delete: [promptId] }) }));
|
||||
if (isRunning) requests.push(fetch(`${comfyBaseUrl}/interrupt`, { method: "POST" }));
|
||||
await Promise.allSettled(requests);
|
||||
}
|
||||
|
||||
function abortableSleep(ms: number, signal?: AbortSignal): Promise<void> {
|
||||
if (!signal) return Bun.sleep(ms);
|
||||
return new Promise((resolve, reject) => {
|
||||
const timeout = setTimeout(resolve, ms);
|
||||
signal.addEventListener("abort", () => {
|
||||
clearTimeout(timeout);
|
||||
reject(signal.reason ?? new DOMException("Generation cancelled", "AbortError"));
|
||||
}, { once: true });
|
||||
});
|
||||
}
|
||||
|
||||
export function selectGeneratedOutputImage(history: unknown): { filename: string; subfolder?: string; type?: string } | undefined {
|
||||
const outputs = (history as { outputs?: Record<string, { images?: { filename: string; subfolder?: string; type?: string }[] }> }).outputs ?? {};
|
||||
const saveImageOutput = outputs["8"]?.images?.find(isGeneratedImage);
|
||||
|
||||
Reference in New Issue
Block a user