diff --git a/custom_routes.py b/custom_routes.py index de0ab6d..54441c9 100644 --- a/custom_routes.py +++ b/custom_routes.py @@ -672,6 +672,14 @@ async def comfy_deploy_run(request): return web.json_response(res, status=status) +@server.PromptServer.instance.routes.post("/comfyui-deploy/interrupt") +async def interrupt_prompt(request): + data = await request.json() + prompt_id = data.get("prompt_id") + await update_run(prompt_id, Status.CANCELLED) + return web.json_response({"message": "Prompt interrupted"}, status=200) + + async def stream_prompt(data, token): # In older version, we use workflow_api, but this has inputs already swapped in nextjs frontend, which is tricky workflow_api = data.get("workflow_api_raw") diff --git a/globals.py b/globals.py index 8588c42..08456ea 100644 --- a/globals.py +++ b/globals.py @@ -18,6 +18,7 @@ class Status(Enum): SUCCESS = "success" FAILED = "failed" UPLOADING = "uploading" + CANCELLED = "cancelled" class StreamingPrompt(BaseModel): diff --git a/web-plugin/index.js b/web-plugin/index.js index b5ad534..3a428bf 100644 --- a/web-plugin/index.js +++ b/web-plugin/index.js @@ -2672,31 +2672,53 @@ const orginal_fetch_api = api.fetchApi; api.fetchApi = async (route, options) => { // console.log("Fetch API called with args:", route, options, ext.native_mode); - if (route.startsWith("/prompt") && ext.native_mode) { - const info = await getSelectedWorkflowInfo(); + if (ext.native_mode) { + if (route.startsWith("/prompt")) { + const info = await getSelectedWorkflowInfo(); - if (!info.workflow_id) { - console.log("No workflow id found, fallback to original fetch"); - return await orginal_fetch_api.call(api, route, options); - } + if (!info.workflow_id) { + console.log("No workflow id found, fallback to original fetch"); + return await orginal_fetch_api.call(api, route, options); + } - console.log("info", info); - if (info) { + console.log("info", info); + if (info) { + const body = JSON.parse(options.body); + + const data = { + client_id: body.client_id, + workflow_api_json: body.prompt, + workflow: body?.extra_data?.extra_pnginfo?.workflow, + is_native_run: true, + machine_id: info.machine_id, + workflow_id: info.workflow_id, + native_run_api_endpoint: info.native_run_api_endpoint, + gpu_event_id: info.gpu_event_id, + gpu: info.gpu, + }; + + return await fetch("/comfyui-deploy/run", { + method: "POST", + headers: { + Authorization: `Bearer ${info.cd_token}`, + "Content-Type": "application/json", + }, + body: JSON.stringify(data), + }); + } + } else if (route.startsWith("/interrupt")) { + const info = await getSelectedWorkflowInfo(); + + if (!info.workflow_id) { + console.log("No workflow id found, fallback to original fetch"); + return await orginal_fetch_api.call(api, route, options); + } const body = JSON.parse(options.body); - const data = { - client_id: body.client_id, - workflow_api_json: body.prompt, - workflow: body?.extra_data?.extra_pnginfo?.workflow, - is_native_run: true, - machine_id: info.machine_id, - workflow_id: info.workflow_id, - native_run_api_endpoint: info.native_run_api_endpoint, - gpu_event_id: info.gpu_event_id, - gpu: info.gpu, + prompt_id: body.prompt_id, }; - - return await fetch("/comfyui-deploy/run", { + const original_response = await orginal_fetch_api.call(api, route, options); + await fetch("/comfyui-deploy/interrupt", { method: "POST", headers: { Authorization: `Bearer ${info.cd_token}`, @@ -2704,6 +2726,7 @@ api.fetchApi = async (route, options) => { }, body: JSON.stringify(data), }); + return original_response; } }