chunzhi/comfyui-deploy
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add create workflow run error server action catch, redirect workflow parse error, add key revoked col

Browse Source
This commit is contained in:
BennyKok 2023-12-16 14:55:30 +08:00
parent 0cf6d97f2f
commit 21a17cb753
16 changed files with 839 additions and 188 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
custom_routes.py
View File

@ -14,7 +14,7 @@ import io
import time
import execution
import random
import traceback
import uuid
import asyncio
import atexit
@ -90,7 +90,14 @@ async def comfy_deploy_run(request):
"client_id": "fake_client" #api.client_id
}
res = post_prompt(prompt)
try:
res = post_prompt(prompt)
except Exception as e:
error_type = type(e).__name__
stack_trace = traceback.format_exc().strip().split('\n')[-2]
print(f"error: {error_type}, {e}")
print(f"stack trace: {stack_trace}")
return web.Response(status=500, reason=f"{error_type}: {e}, {stack_trace}")
prompt_metadata[res['prompt_id']] = {
'status_endpoint': data.get('status_endpoint'),
@ -252,5 +259,9 @@ async def update_run_with_output(prompt_id, data):
}
requests.post(status_endpoint, json=body)
await send('outputs_uploaded', {
"prompt_id": prompt_id
})
prompt_server.send_json_original = prompt_server.send_json
prompt_server.send_json = send_json_override.__get__(prompt_server, server.PromptServer)

8
web-plugin/index.js
View File

@ -5,7 +5,7 @@ import { ComfyWidgets, LGraphNode } from "./widgets.js";
/** @typedef {import('../../../web/types/comfy.js').ComfyExtension} ComfyExtension*/
/** @type {ComfyExtension} */
const ext = {
name: "BennyKok.ComfyDeploy",
name: "BennyKok.ComfyUIDeploy",
init(app) {
addButton();
@ -65,7 +65,7 @@ const ext = {
// Load default visibility
LiteGraph.registerNodeType(
"Comfy Deploy",
"ComfyDeploy",
Object.assign(ComfyDeploy, {
title_mode: LiteGraph.NORMAL_TITLE,
title: "Comfy Deploy",
@ -112,7 +112,7 @@ function addButton() {
/** @type {LGraph} */
const graph = app.graph;
const deployMeta = graph.findNodesByType("Comfy Deploy");
const deployMeta = graph.findNodesByType("ComfyDeploy");
const deployMetaNode = deployMeta[0];
console.log(deployMetaNode);
@ -137,6 +137,8 @@ function addButton() {
deploy.textContent = "Deploying...";
deploy.style.color = "orange";
console.log(prompt);
const apiRoute = endpoint + "/api/upload"
// const userId = apiKey
try {

1
web/drizzle/0010_organic_marten_broadcloak.sql Normal file
View File

@ -0,0 +1 @@
ALTER TABLE "comfyui_deploy"."api_keys" ADD COLUMN "revoked" boolean DEFAULT false NOT NULL;

621
web/drizzle/meta/0010_snapshot.json Normal file
View File

@ -0,0 +1,621 @@
{
"id": "128aae04-697c-4253-b7de-b27e895e03f8",
"prevId": "a5d542ea-484f-431a-b31f-170e789a03a7",
"version": "5",
"dialect": "pg",
"tables": {
"api_keys": {
"name": "api_keys",
"schema": "comfyui_deploy",
"columns": {
"id": {
"name": "id",
"type": "uuid",
"primaryKey": true,
"notNull": true,
"default": "gen_random_uuid()"
},
"key": {
"name": "key",
"type": "text",
"primaryKey": false,
"notNull": true
},
"name": {
"name": "name",
"type": "text",
"primaryKey": false,
"notNull": true
},
"user_id": {
"name": "user_id",
"type": "text",
"primaryKey": false,
"notNull": true
},
"org_id": {
"name": "org_id",
"type": "text",
"primaryKey": false,
"notNull": false
},
"revoked": {
"name": "revoked",
"type": "boolean",
"primaryKey": false,
"notNull": true,
"default": false
},
"created_at": {
"name": "created_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
},
"updated_at": {
"name": "updated_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
}
},
"indexes": {},
"foreignKeys": {
"api_keys_user_id_users_id_fk": {
"name": "api_keys_user_id_users_id_fk",
"tableFrom": "api_keys",
"tableTo": "users",
"columnsFrom": [
"user_id"
],
"columnsTo": [
"id"
],
"onDelete": "cascade",
"onUpdate": "no action"
}
},
"compositePrimaryKeys": {},
"uniqueConstraints": {
"api_keys_key_unique": {
"name": "api_keys_key_unique",
"nullsNotDistinct": false,
"columns": [
"key"
]
}
}
},
"deployments": {
"name": "deployments",
"schema": "comfyui_deploy",
"columns": {
"id": {
"name": "id",
"type": "uuid",
"primaryKey": true,
"notNull": true,
"default": "gen_random_uuid()"
},
"user_id": {
"name": "user_id",
"type": "text",
"primaryKey": false,
"notNull": true
},
"workflow_version_id": {
"name": "workflow_version_id",
"type": "uuid",
"primaryKey": false,
"notNull": true
},
"workflow_id": {
"name": "workflow_id",
"type": "uuid",
"primaryKey": false,
"notNull": true
},
"machine_id": {
"name": "machine_id",
"type": "uuid",
"primaryKey": false,
"notNull": true
},
"environment": {
"name": "environment",
"type": "deployment_environment",
"primaryKey": false,
"notNull": true
},
"created_at": {
"name": "created_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
},
"updated_at": {
"name": "updated_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
}
},
"indexes": {},
"foreignKeys": {
"deployments_user_id_users_id_fk": {
"name": "deployments_user_id_users_id_fk",
"tableFrom": "deployments",
"tableTo": "users",
"columnsFrom": [
"user_id"
],
"columnsTo": [
"id"
],
"onDelete": "cascade",
"onUpdate": "no action"
},
"deployments_workflow_version_id_workflow_versions_id_fk": {
"name": "deployments_workflow_version_id_workflow_versions_id_fk",
"tableFrom": "deployments",
"tableTo": "workflow_versions",
"columnsFrom": [
"workflow_version_id"
],
"columnsTo": [
"id"
],
"onDelete": "no action",
"onUpdate": "no action"
},
"deployments_workflow_id_workflows_id_fk": {
"name": "deployments_workflow_id_workflows_id_fk",
"tableFrom": "deployments",
"tableTo": "workflows",
"columnsFrom": [
"workflow_id"
],
"columnsTo": [
"id"
],
"onDelete": "no action",
"onUpdate": "no action"
},
"deployments_machine_id_machines_id_fk": {
"name": "deployments_machine_id_machines_id_fk",
"tableFrom": "deployments",
"tableTo": "machines",
"columnsFrom": [
"machine_id"
],
"columnsTo": [
"id"
],
"onDelete": "no action",
"onUpdate": "no action"
}
},
"compositePrimaryKeys": {},
"uniqueConstraints": {}
},
"machines": {
"name": "machines",
"schema": "comfyui_deploy",
"columns": {
"id": {
"name": "id",
"type": "uuid",
"primaryKey": true,
"notNull": true,
"default": "gen_random_uuid()"
},
"user_id": {
"name": "user_id",
"type": "text",
"primaryKey": false,
"notNull": true
},
"name": {
"name": "name",
"type": "text",
"primaryKey": false,
"notNull": true
},
"endpoint": {
"name": "endpoint",
"type": "text",
"primaryKey": false,
"notNull": true
},
"created_at": {
"name": "created_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
},
"updated_at": {
"name": "updated_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
}
},
"indexes": {},
"foreignKeys": {
"machines_user_id_users_id_fk": {
"name": "machines_user_id_users_id_fk",
"tableFrom": "machines",
"tableTo": "users",
"columnsFrom": [
"user_id"
],
"columnsTo": [
"id"
],
"onDelete": "cascade",
"onUpdate": "no action"
}
},
"compositePrimaryKeys": {},
"uniqueConstraints": {}
},
"users": {
"name": "users",
"schema": "comfyui_deploy",
"columns": {
"id": {
"name": "id",
"type": "text",
"primaryKey": true,
"notNull": true
},
"username": {
"name": "username",
"type": "text",
"primaryKey": false,
"notNull": true
},
"name": {
"name": "name",
"type": "text",
"primaryKey": false,
"notNull": true
},
"created_at": {
"name": "created_at",
"type": "timestamp",
"primaryKey": false,
"notNull": false,
"default": "now()"
},
"updated_at": {
"name": "updated_at",
"type": "timestamp",
"primaryKey": false,
"notNull": false,
"default": "now()"
}
},
"indexes": {},
"foreignKeys": {},
"compositePrimaryKeys": {},
"uniqueConstraints": {}
},
"workflow_run_outputs": {
"name": "workflow_run_outputs",
"schema": "comfyui_deploy",
"columns": {
"id": {
"name": "id",
"type": "uuid",
"primaryKey": true,
"notNull": true,
"default": "gen_random_uuid()"
},
"run_id": {
"name": "run_id",
"type": "uuid",
"primaryKey": false,
"notNull": true
},
"data": {
"name": "data",
"type": "jsonb",
"primaryKey": false,
"notNull": false
},
"created_at": {
"name": "created_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
},
"updated_at": {
"name": "updated_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
}
},
"indexes": {},
"foreignKeys": {
"workflow_run_outputs_run_id_workflow_runs_id_fk": {
"name": "workflow_run_outputs_run_id_workflow_runs_id_fk",
"tableFrom": "workflow_run_outputs",
"tableTo": "workflow_runs",
"columnsFrom": [
"run_id"
],
"columnsTo": [
"id"
],
"onDelete": "cascade",
"onUpdate": "no action"
}
},
"compositePrimaryKeys": {},
"uniqueConstraints": {}
},
"workflow_runs": {
"name": "workflow_runs",
"schema": "comfyui_deploy",
"columns": {
"id": {
"name": "id",
"type": "uuid",
"primaryKey": true,
"notNull": true,
"default": "gen_random_uuid()"
},
"workflow_version_id": {
"name": "workflow_version_id",
"type": "uuid",
"primaryKey": false,
"notNull": false
},
"workflow_inputs": {
"name": "workflow_inputs",
"type": "jsonb",
"primaryKey": false,
"notNull": false
},
"workflow_id": {
"name": "workflow_id",
"type": "uuid",
"primaryKey": false,
"notNull": true
},
"machine_id": {
"name": "machine_id",
"type": "uuid",
"primaryKey": false,
"notNull": false
},
"status": {
"name": "status",
"type": "workflow_run_status",
"primaryKey": false,
"notNull": true,
"default": "'not-started'"
},
"ended_at": {
"name": "ended_at",
"type": "timestamp",
"primaryKey": false,
"notNull": false
},
"created_at": {
"name": "created_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
}
},
"indexes": {},
"foreignKeys": {
"workflow_runs_workflow_version_id_workflow_versions_id_fk": {
"name": "workflow_runs_workflow_version_id_workflow_versions_id_fk",
"tableFrom": "workflow_runs",
"tableTo": "workflow_versions",
"columnsFrom": [
"workflow_version_id"
],
"columnsTo": [
"id"
],
"onDelete": "set null",
"onUpdate": "no action"
},
"workflow_runs_workflow_id_workflows_id_fk": {
"name": "workflow_runs_workflow_id_workflows_id_fk",
"tableFrom": "workflow_runs",
"tableTo": "workflows",
"columnsFrom": [
"workflow_id"
],
"columnsTo": [
"id"
],
"onDelete": "cascade",
"onUpdate": "no action"
},
"workflow_runs_machine_id_machines_id_fk": {
"name": "workflow_runs_machine_id_machines_id_fk",
"tableFrom": "workflow_runs",
"tableTo": "machines",
"columnsFrom": [
"machine_id"
],
"columnsTo": [
"id"
],
"onDelete": "set null",
"onUpdate": "no action"
}
},
"compositePrimaryKeys": {},
"uniqueConstraints": {}
},
"workflows": {
"name": "workflows",
"schema": "comfyui_deploy",
"columns": {
"id": {
"name": "id",
"type": "uuid",
"primaryKey": true,
"notNull": true,
"default": "gen_random_uuid()"
},
"user_id": {
"name": "user_id",
"type": "text",
"primaryKey": false,
"notNull": true
},
"name": {
"name": "name",
"type": "text",
"primaryKey": false,
"notNull": true
},
"created_at": {
"name": "created_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
},
"updated_at": {
"name": "updated_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
}
},
"indexes": {},
"foreignKeys": {
"workflows_user_id_users_id_fk": {
"name": "workflows_user_id_users_id_fk",
"tableFrom": "workflows",
"tableTo": "users",
"columnsFrom": [
"user_id"
],
"columnsTo": [
"id"
],
"onDelete": "cascade",
"onUpdate": "no action"
}
},
"compositePrimaryKeys": {},
"uniqueConstraints": {}
},
"workflow_versions": {
"name": "workflow_versions",
"schema": "comfyui_deploy",
"columns": {
"workflow_id": {
"name": "workflow_id",
"type": "uuid",
"primaryKey": false,
"notNull": true
},
"id": {
"name": "id",
"type": "uuid",
"primaryKey": true,
"notNull": true,
"default": "gen_random_uuid()"
},
"workflow": {
"name": "workflow",
"type": "jsonb",
"primaryKey": false,
"notNull": false
},
"workflow_api": {
"name": "workflow_api",
"type": "jsonb",
"primaryKey": false,
"notNull": false
},
"version": {
"name": "version",
"type": "integer",
"primaryKey": false,
"notNull": true
},
"created_at": {
"name": "created_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
},
"updated_at": {
"name": "updated_at",
"type": "timestamp",
"primaryKey": false,
"notNull": true,
"default": "now()"
}
},
"indexes": {},
"foreignKeys": {
"workflow_versions_workflow_id_workflows_id_fk": {
"name": "workflow_versions_workflow_id_workflows_id_fk",
"tableFrom": "workflow_versions",
"tableTo": "workflows",
"columnsFrom": [
"workflow_id"
],
"columnsTo": [
"id"
],
"onDelete": "cascade",
"onUpdate": "no action"
}
},
"compositePrimaryKeys": {},
"uniqueConstraints": {}
}
},
"enums": {
"deployment_environment": {
"name": "deployment_environment",
"values": {
"staging": "staging",
"production": "production"
}
},
"workflow_run_status": {
"name": "workflow_run_status",
"values": {
"not-started": "not-started",
"running": "running",
"success": "success",
"failed": "failed"
}
}
},
"schemas": {
"comfyui_deploy": "comfyui_deploy"
},
"_meta": {
"schemas": {},
"tables": {},
"columns": {}
}
}

7
web/drizzle/meta/_journal.json
View File

@ -71,6 +71,13 @@
"when": 1702632471725,
"tag": "0009_easy_banshee",
"breakpoints": true
},
{
"idx": 10,
"version": "5",
"when": 1702708886828,
"tag": "0010_organic_marten_broadcloak",
"breakpoints": true
}
]
}

60
web/src/app/api/create-run/route.ts
View File

@ -1,60 +0,0 @@
import { parseDataSafe } from "../../../lib/parseDataSafe";
import { createRun } from "../../../server/createRun";
import { getRunsOutput } from "@/server/getRunsOutput";
import { NextResponse } from "next/server";
import { z } from "zod";
const Request = z.object({
workflow_version_id: z.string(),
machine_id: z.string(),
});
const Request2 = z.object({
run_id: z.string(),
});
export async function GET(request: Request) {
const [data, error] = await parseDataSafe(Request2, request);
if (!data || error) return error;
const run = await getRunsOutput(data.run_id);
return NextResponse.json(run, {
status: 200,
});
}
export async function POST(request: Request) {
const [data, error] = await parseDataSafe(Request, request);
if (!data || error) return error;
const origin = new URL(request.url).origin;
const { workflow_version_id, machine_id } = data;
try {
const workflow_run_id = await createRun(
origin,
workflow_version_id,
machine_id
);
return NextResponse.json(
{
workflow_run_id: workflow_run_id.workflow_run_id,
},
{
status: 200,
}
);
} catch (error: any) {
return NextResponse.json(
{
error: error.message,
},
{
status: 500,
}
);
}
}

34
web/src/app/api/run/route.ts
View File

@ -2,6 +2,7 @@ import { parseDataSafe } from "../../../lib/parseDataSafe";
import { createRun } from "../../../server/createRun";
import { db } from "@/db/db";
import { deploymentsTable } from "@/db/schema";
import { isKeyRevoked } from "@/server/curdApiKeys";
import { getRunsData } from "@/server/getRunsOutput";
import { parseJWT } from "@/server/parseJWT";
import { replaceCDNUrl } from "@/server/resource";
@ -18,14 +19,26 @@ const Request2 = z.object({
run_id: z.string(),
});
export async function GET(request: Request) {
async function checkToken(request: Request) {
const token = request.headers.get("Authorization")?.split(" ")?.[1]; // Assuming token is sent as "Bearer your_token"
const userData = token ? parseJWT(token) : undefined;
if (!userData) {
if (!userData || token === undefined) {
return new NextResponse("Invalid or expired token", {
status: 401,
});
} else {
const revokedKey = await isKeyRevoked(token);
if (revokedKey)
return new NextResponse("Revoked token", {
status: 401,
});
}
}
export async function GET(request: Request) {
const invalidRequest = await checkToken(request)
if (invalidRequest) return invalidRequest;
const [data, error] = await parseDataSafe(Request2, request);
if (!data || error) return error;
@ -41,7 +54,7 @@ export async function GET(request: Request) {
for (let j = 0; j < output.data?.images.length; j++) {
const element = output.data?.images[j];
element.url = replaceCDNUrl(
`${process.env.SPACES_ENDPOINT}/${process.env.SPACES_BUCKET}/outputs/runs/${run.id}/${element.filename}`
`${process.env.SPACES_ENDPOINT}/${process.env.SPACES_BUCKET}/outputs/runs/${run.id}/${element.filename}`,
);
}
}
@ -53,13 +66,8 @@ export async function GET(request: Request) {
}
export async function POST(request: Request) {
const token = request.headers.get("Authorization")?.split(" ")?.[1]; // Assuming token is sent as "Bearer your_token"
const userData = token ? parseJWT(token) : undefined;
if (!userData) {
return new NextResponse("Invalid or expired token", {
status: 401,
});
}
const invalidRequest = await checkToken(request)
if (invalidRequest) return invalidRequest;
const [data, error] = await parseDataSafe(Request, request);
if (!data || error) return error;
@ -79,7 +87,7 @@ export async function POST(request: Request) {
origin,
deploymentData.workflow_version_id,
deploymentData.machine_id,
inputs
inputs,
);
return NextResponse.json(
@ -88,7 +96,7 @@ export async function POST(request: Request) {
},
{
status: 200,
}
},
);
} catch (error: any) {
return NextResponse.json(
@ -97,7 +105,7 @@ export async function POST(request: Request) {
},
{
status: 500,
}
},
);
}
}

49
web/src/app/layout.tsx
View File

@ -1,5 +1,8 @@
import "./globals.css";
import { NavbarRight } from "@/components/NavbarRight";
import { Button } from "@/components/ui/button";
import { ClerkProvider, UserButton } from "@clerk/nextjs";
import { Github } from "lucide-react";
import type { Metadata } from "next";
import { Inter } from "next/font/google";
import { Toaster } from "sonner";
@ -18,21 +21,37 @@ export default function RootLayout({
}) {
return (
<html lang="en">
<body className={inter.className}>
<main className="flex min-h-screen flex-col items-center justify-start">
<div className="w-full h-18 flex items-center gap-4 p-4 border-b border-gray-200">
<a className="font-bold text-lg hover:underline" href="/">
Comfy Deploy
</a>
<NavbarRight />
{/* <div></div> */}
</div>
<div className="md:px-10 px-6 w-full flex items-start">
{children}
</div>
<Toaster richColors />
</main>
</body>
<ClerkProvider>
<body className={inter.className}>
<main className="flex min-h-screen flex-col items-center justify-start">
<div className="w-full h-18 flex items-center justify-between gap-4 p-4 border-b border-gray-200">
<div className="flex flex-row items-center gap-4">
<a className="font-bold text-lg hover:underline" href="/">
ComfyUI Deploy
</a>
<NavbarRight />
</div>
<div className="flex flex-row items-center gap-2">
<UserButton />
<Button
asChild
variant={"outline"}
className="rounded-full aspect-square p-2"
>
<a target="_blank" href="https://github.com/BennyKok/comfyui-deploy">
<Github />
</a>
</Button>
</div>
{/* <div></div> */}
</div>
<div className="md:px-10 px-6 w-full flex items-start">
{children}
</div>
<Toaster richColors />
</main>
</body>
</ClerkProvider>
</html>
);
}

12
web/src/components/LiveStatus.tsx
View File

@ -4,6 +4,8 @@ import { useStore } from "@/components/MachinesWS";
import { StatusBadge } from "@/components/StatusBadge";
import { TableCell } from "@/components/ui/table";
import { type findAllRuns } from "@/server/findAllRuns";
import { useRouter } from "next/navigation";
import { useEffect, useState } from "react";
export function LiveStatus({
run,
@ -14,7 +16,7 @@ export function LiveStatus({
(state) =>
state.data
.filter((x) => x.id === run.id)
.sort((a, b) => b.timestamp - a.timestamp)?.[0]
.sort((a, b) => b.timestamp - a.timestamp)?.[0],
);
let status = run.status;
@ -26,6 +28,14 @@ export function LiveStatus({
status = "running";
}
const router = useRouter();
useEffect(() => {
if (data?.json.event === "outputs_uploaded") {
router.refresh()
}
}, [data?.json.event]);
return (
<>
<TableCell>

3
web/src/components/MachinesWS.tsx
View File

@ -55,7 +55,8 @@ function MachineWS({
const { lastMessage, readyState } = useWebSocket(
`${wsEndpoint}/comfyui-deploy/ws`,
{
reconnectAttempts: 10,
shouldReconnect: ()=> true,
reconnectAttempts: 20,
reconnectInterval: 1000,
}
);

2
web/src/components/RunsTable.tsx
View File

@ -38,7 +38,7 @@ export async function RunsTable(props: { workflow_id: string }) {
export async function DeploymentsTable(props: { workflow_id: string }) {
const allRuns = await findAllDeployments(props.workflow_id);
return (
<div className="overflow-auto h-[400px] w-full">
<div className="overflow-auto h-fit lg:h-[400px] w-full">
<Table className="">
<TableCaption>A list of your deployments</TableCaption>
<TableHeader className="bg-background top-0 sticky">

2
web/src/components/callServerPromise.tsx
View File

@ -7,6 +7,8 @@ export async function callServerPromise<T>(result: Promise<T>) {
.then((x) => {
if ((x as { message: string })?.message !== undefined) {
toast.success((x as { message: string }).message);
} else if ((x as { error: string })?.error !== undefined) {
toast.error((x as { error: string }).error);
}
return x;
})

4
web/src/db/schema.ts
View File

@ -7,6 +7,7 @@ import {
timestamp,
jsonb,
pgEnum,
boolean,
} from "drizzle-orm/pg-core";
import { z } from "zod";
@ -52,7 +53,7 @@ export const workflowType = z.any();
export const workflowAPIType = z.record(
z.object({
inputs: z.record(z.any()),
class_type: z.string(),
class_type: z.string().optional(),
})
);
@ -213,6 +214,7 @@ export const apiKeyTable = dbSchema.table("api_keys", {
})
.notNull(),
org_id: text("org_id"),
revoked: boolean("revoked").default(false).notNull(),
created_at: timestamp("created_at").defaultNow().notNull(),
updated_at: timestamp("updated_at").defaultNow().notNull(),
});

171
web/src/server/createRun.ts
View File

@ -6,102 +6,101 @@ import { ComfyAPI_Run } from "@/types/ComfyAPI_Run";
import { eq } from "drizzle-orm";
import { revalidatePath } from "next/cache";
import "server-only";
import { withServerPromise } from "./withServerPromise";
export async function createRun(
origin: string,
workflow_version_id: string,
machine_id: string,
inputs?: Record<string, string>
) {
const machine = await db.query.machinesTable.findFirst({
where: eq(workflowRunsTable.id, machine_id),
});
if (!machine) {
throw new Error("Machine not found");
// return new Response("Machine not found", {
// status: 404,
// });
}
const workflow_version_data =
// workflow_version_id
// ?
await db.query.workflowVersionTable.findFirst({
where: eq(workflowRunsTable.id, workflow_version_id),
export const createRun = withServerPromise(
async (
origin: string,
workflow_version_id: string,
machine_id: string,
inputs?: Record<string, string>,
) => {
const machine = await db.query.machinesTable.findFirst({
where: eq(workflowRunsTable.id, machine_id),
});
// : workflow_version != undefined
// ? await db.query.workflowVersionTable.findFirst({
// where: and(
// eq(workflowVersionTable.version, workflow_version),
// eq(workflowVersionTable.workflow_id)
// ),
// })
// : null;
if (!workflow_version_data) {
throw new Error("Workflow version not found");
// return new Response("Workflow version not found", {
// status: 404,
// });
}
const comfyui_endpoint = `${machine.endpoint}/comfyui-deploy/run`;
const workflow_api = workflow_version_data.workflow_api;
// Replace the inputs
if (inputs && workflow_api) {
for (const key in inputs) {
Object.entries(workflow_api).forEach(([_, node]) => {
if (node.inputs["input_id"] === key) {
node.inputs["input_id"] = inputs[key];
}
});
if (!machine) {
throw new Error("Machine not found");
// return new Response("Machine not found", {
// status: 404,
// });
}
}
const body = {
workflow_api: workflow_api,
status_endpoint: `${origin}/api/update-run`,
file_upload_endpoint: `${origin}/api/file-upload`,
};
// console.log(body);
const bodyJson = JSON.stringify(body);
// console.log(bodyJson);
const workflow_version_data =
// workflow_version_id
// ?
await db.query.workflowVersionTable.findFirst({
where: eq(workflowRunsTable.id, workflow_version_id),
});
// : workflow_version != undefined
// ? await db.query.workflowVersionTable.findFirst({
// where: and(
// eq(workflowVersionTable.version, workflow_version),
// eq(workflowVersionTable.workflow_id)
// ),
// })
// : null;
if (!workflow_version_data) {
throw new Error("Workflow version not found");
// return new Response("Workflow version not found", {
// status: 404,
// });
}
// Sending to comfyui
const _result = await fetch(comfyui_endpoint, {
method: "POST",
body: bodyJson,
cache: "no-store",
});
const comfyui_endpoint = `${machine.endpoint}/comfyui-deploy/run`;
if (!_result.ok) {
throw new Error(`Error creating run, ${_result.statusText}`);
}
const workflow_api = workflow_version_data.workflow_api;
console.log(_result);
// Replace the inputs
if (inputs && workflow_api) {
for (const key in inputs) {
Object.entries(workflow_api).forEach(([_, node]) => {
if (node.inputs["input_id"] === key) {
node.inputs["input_id"] = inputs[key];
}
});
}
}
const result = await ComfyAPI_Run.parseAsync(await _result.json());
const body = {
workflow_api: workflow_api,
status_endpoint: `${origin}/api/update-run`,
file_upload_endpoint: `${origin}/api/file-upload`,
};
// console.log(body);
const bodyJson = JSON.stringify(body);
// console.log(bodyJson);
console.log(result);
// Sending to comfyui
const _result = await fetch(comfyui_endpoint, {
method: "POST",
body: bodyJson,
cache: "no-store",
});
// Add to our db
const workflow_run = await db
.insert(workflowRunsTable)
.values({
id: result.prompt_id,
workflow_id: workflow_version_data.workflow_id,
workflow_version_id: workflow_version_data.id,
workflow_inputs: inputs,
machine_id,
})
.returning();
if (!_result.ok) {
throw new Error(`Error creating run, ${_result.statusText}`);
}
revalidatePath(`/${workflow_version_data.workflow_id}`);
const result = await ComfyAPI_Run.parseAsync(await _result.json());
return {
workflow_run_id: workflow_run[0].id,
message: "Successfully workflow run",
};
}
// Add to our db
const workflow_run = await db
.insert(workflowRunsTable)
.values({
id: result.prompt_id,
workflow_id: workflow_version_data.workflow_id,
workflow_version_id: workflow_version_data.id,
workflow_inputs: inputs,
machine_id,
})
.returning();
revalidatePath(`/${workflow_version_data.workflow_id}`);
return {
workflow_run_id: workflow_run[0].id,
message: "Successfully workflow run",
};
},
);

26
web/src/server/curdApiKeys.ts
View File

@ -29,7 +29,7 @@ export async function addNewAPIKey(name: string) {
if (orgId) {
token = jwt.sign(
{ user_id: userId, org_id: orgId },
process.env.JWT_SECRET!
process.env.JWT_SECRET!,
);
} else {
token = jwt.sign({ user_id: userId }, process.env.JWT_SECRET!);
@ -57,12 +57,20 @@ export async function deleteAPIKey(id: string) {
if (orgId) {
await db
.delete(apiKeyTable)
.update(apiKeyTable)
.set({
revoked: true,
updated_at: new Date(),
})
.where(and(eq(apiKeyTable.id, id), eq(apiKeyTable.org_id, orgId)))
.execute();
} else {
await db
.delete(apiKeyTable)
.update(apiKeyTable)
.set({
revoked: true,
updated_at: new Date(),
})
.where(and(eq(apiKeyTable.id, id), eq(apiKeyTable.user_id, userId)))
.execute();
}
@ -77,13 +85,21 @@ export async function getAPIKeys() {
if (orgId) {
return await db.query.apiKeyTable.findMany({
where: eq(apiKeyTable.org_id, orgId),
where: and(eq(apiKeyTable.org_id, orgId), eq(apiKeyTable.revoked, false)),
orderBy: desc(apiKeyTable.created_at),
});
} else {
return await db.query.apiKeyTable.findMany({
where: eq(apiKeyTable.user_id, userId),
where: and(eq(apiKeyTable.user_id, userId), eq(apiKeyTable.revoked, false)),
orderBy: desc(apiKeyTable.created_at),
});
}
}
export async function isKeyRevoked(key: string) {
const revokedKey = await db.query.apiKeyTable.findFirst({
where: and(eq(apiKeyTable.key, key), eq(apiKeyTable.revoked, true)),
});
return revokedKey !== undefined;
}

12
web/src/server/withServerPromise.ts Normal file
View File

@ -0,0 +1,12 @@
export async function wrapServerPromise<T>(result: Promise<T>) {
return result.catch((error) => {
return {
error: error.message,
};
});
}
export function withServerPromise<T extends (...args: any[]) => Promise<any>>(
fn: T
): (...args: Parameters<T>) => Promise<ReturnType<T> | { error: string; }> {
return (...args: Parameters<T>) => wrapServerPromise(fn(...args));
}