diff --git a/custom_routes.py b/custom_routes.py index ab047b2..b120033 100644 --- a/custom_routes.py +++ b/custom_routes.py @@ -51,7 +51,8 @@ def post_prompt(json_data): if "client_id" in json_data: extra_data["client_id"] = json_data["client_id"] if valid[0]: - prompt_id = str(uuid.uuid4()) + # if the prompt id is provided + prompt_id = json_data.get("prompt_id") or str(uuid.uuid4()) outputs_to_execute = valid[2] prompt_server.prompt_queue.put( (number, prompt_id, prompt, extra_data, outputs_to_execute) @@ -80,13 +81,17 @@ async def comfy_deploy_run(request): workflow_api = data.get("workflow_api") + # The prompt id generated from comfy deploy, can be None + prompt_id = data.get("prompt_id") + for key in workflow_api: if 'inputs' in workflow_api[key] and 'seed' in workflow_api[key]['inputs']: workflow_api[key]['inputs']['seed'] = randomSeed() prompt = { "prompt": workflow_api, - "client_id": "comfy_deploy_instance" #api.client_id + "client_id": "comfy_deploy_instance", #api.client_id + "prompt_id": prompt_id } try: @@ -136,6 +141,19 @@ async def websocket_handler(request): sockets.pop(sid, None) return ws +@server.PromptServer.instance.routes.get('/comfyui-deploy/check-status') +async def comfy_deploy_check_status(request): + prompt_server = server.PromptServer.instance + prompt_id = request.rel_url.query.get('prompt_id', None) + if prompt_id in prompt_metadata: + return web.json_response({ + "status": prompt_metadata[prompt_id]['status'].value + }) + else: + return web.json_response({ + "message": "prompt_id not found" + }) + async def send(event, data, sid=None): try: if sid: @@ -168,7 +186,7 @@ async def send_json_override(self, event, data, sid=None): update_run(prompt_id, Status.RUNNING) # the last executing event is none, then the workflow is finished - if event == 'executing' and data.get('node') is None: + if event == 'executing' and data.get('node') is None and not have_pending_upload(prompt_id): update_run(prompt_id, Status.SUCCESS) if event == 'execution_error': @@ -176,7 +194,7 @@ async def send_json_override(self, event, data, sid=None): asyncio.create_task(update_run_with_output(prompt_id, data)) if event == 'executed' and 'node' in data and 'output' in data: - asyncio.create_task(update_run_with_output(prompt_id, data.get('output'))) + asyncio.create_task(update_run_with_output(prompt_id, data.get('output'), node_id=data.get('node'))) # update_run_with_output(prompt_id, data.get('output')) @@ -185,6 +203,7 @@ class Status(Enum): RUNNING = "running" SUCCESS = "success" FAILED = "failed" + UPLOADING = "uploading" def update_run(prompt_id, status: Status): if prompt_id not in prompt_metadata: @@ -260,7 +279,47 @@ async def upload_file(prompt_id, filename, subfolder=None, content_type="image/p response = requests.put(ok.get("url"), headers=headers, data=data) print("upload file response", response.status_code) -async def update_run_with_output(prompt_id, data): +def have_pending_upload(prompt_id): + if 'uploading_nodes' in prompt_metadata[prompt_id] and len(prompt_metadata[prompt_id]['uploading_nodes']) > 0: + return True + return False + +async def update_file_status(prompt_id, data, uploading, have_error=False, node_id=None): + if 'uploading_nodes' not in prompt_metadata[prompt_id]: + prompt_metadata[prompt_id]['uploading_nodes'] = set() + + if node_id is not None: + if uploading: + prompt_metadata[prompt_id]['uploading_nodes'].add(node_id) + else: + prompt_metadata[prompt_id]['uploading_nodes'].discard(node_id) + + # print(prompt_metadata[prompt_id]) + # Update the remote status + + if have_error: + update_run(prompt_id, Status.FAILED) + await send("failed", { + "prompt_id": prompt_id, + }) + return + + # if there are still nodes that are uploading, then we set the status to uploading + if uploading and have_pending_upload(prompt_id): + if prompt_metadata[prompt_id]['status'] != Status.UPLOADING: + update_run(prompt_id, Status.UPLOADING) + await send("uploading", { + "prompt_id": prompt_id, + }) + + # if there are no nodes that are uploading, then we set the status to success + elif not uploading: + update_run(prompt_id, Status.SUCCESS) + await send("success", { + "prompt_id": prompt_id, + }) + +async def update_run_with_output(prompt_id, data, node_id=None): if prompt_id in prompt_metadata: status_endpoint = prompt_metadata[prompt_id]['status_endpoint'] @@ -270,6 +329,13 @@ async def update_run_with_output(prompt_id, data): } try: + have_upload = 'images' in data or 'files' in data + + print("have_upload", have_upload) + + if have_upload: + await update_file_status(prompt_id, data, True, node_id=node_id) + images = data.get('images', []) for image in images: await upload_file(prompt_id, image.get("filename"), subfolder=image.get("subfolder"), type=image.get("type"), content_type=image.get("content_type", "image/png")) @@ -278,6 +344,9 @@ async def update_run_with_output(prompt_id, data): for file in files: await upload_file(prompt_id, file.get("filename"), subfolder=file.get("subfolder"), type=file.get("type"), content_type=image.get("content_type", "image/png")) + if have_upload: + await update_file_status(prompt_id, data, False, node_id=node_id) + except Exception as e: error_type = type(e).__name__ stack_trace = traceback.format_exc().strip() @@ -291,6 +360,7 @@ async def update_run_with_output(prompt_id, data): } } } + await update_file_status(prompt_id, data, False, have_error=True) print(body) print(f"Error occurred while uploading file: {e}") diff --git a/web/bun.lockb b/web/bun.lockb index a803782..704c778 100755 Binary files a/web/bun.lockb and b/web/bun.lockb differ diff --git a/web/drizzle/0014_short_sunfire.sql b/web/drizzle/0014_short_sunfire.sql new file mode 100644 index 0000000..1f69259 --- /dev/null +++ b/web/drizzle/0014_short_sunfire.sql @@ -0,0 +1,7 @@ +DO $$ BEGIN + CREATE TYPE "machine_type" AS ENUM('classic', 'runpod-serverless'); +EXCEPTION + WHEN duplicate_object THEN null; +END $$; +--> statement-breakpoint +ALTER TABLE "comfyui_deploy"."machines" ADD COLUMN "type" "machine_type" DEFAULT 'classic' NOT NULL; \ No newline at end of file diff --git a/web/drizzle/0015_simple_killmonger.sql b/web/drizzle/0015_simple_killmonger.sql new file mode 100644 index 0000000..97662e7 --- /dev/null +++ b/web/drizzle/0015_simple_killmonger.sql @@ -0,0 +1 @@ +ALTER TABLE "comfyui_deploy"."machines" ADD COLUMN "auth_token" text; \ No newline at end of file diff --git a/web/drizzle/meta/0014_snapshot.json b/web/drizzle/meta/0014_snapshot.json new file mode 100644 index 0000000..0b62c5b --- /dev/null +++ b/web/drizzle/meta/0014_snapshot.json @@ -0,0 +1,657 @@ +{ + "id": "f29ccb7c-8e75-403f-bef6-60758f9db7c7", + "prevId": "61c93b0c-9eae-46b9-bed9-26ef9eed4d19", + "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()" + }, + "disabled": { + "name": "disabled", + "type": "boolean", + "primaryKey": false, + "notNull": true, + "default": false + }, + "type": { + "name": "type", + "type": "machine_type", + "primaryKey": false, + "notNull": true, + "default": "'classic'" + } + }, + "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 + }, + "origin": { + "name": "origin", + "type": "workflow_run_origin", + "primaryKey": false, + "notNull": true, + "default": "'api'" + }, + "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" + } + }, + "machine_type": { + "name": "machine_type", + "values": { + "classic": "classic", + "runpod-serverless": "runpod-serverless" + } + }, + "workflow_run_origin": { + "name": "workflow_run_origin", + "values": { + "manual": "manual", + "api": "api" + } + }, + "workflow_run_status": { + "name": "workflow_run_status", + "values": { + "not-started": "not-started", + "running": "running", + "uploading": "uploading", + "success": "success", + "failed": "failed" + } + } + }, + "schemas": { + "comfyui_deploy": "comfyui_deploy" + }, + "_meta": { + "schemas": {}, + "tables": {}, + "columns": {} + } +} \ No newline at end of file diff --git a/web/drizzle/meta/0015_snapshot.json b/web/drizzle/meta/0015_snapshot.json new file mode 100644 index 0000000..dba1c20 --- /dev/null +++ b/web/drizzle/meta/0015_snapshot.json @@ -0,0 +1,663 @@ +{ + "id": "92bef822-0089-48aa-8f43-5bba40cdce2e", + "prevId": "f29ccb7c-8e75-403f-bef6-60758f9db7c7", + "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()" + }, + "disabled": { + "name": "disabled", + "type": "boolean", + "primaryKey": false, + "notNull": true, + "default": false + }, + "auth_token": { + "name": "auth_token", + "type": "text", + "primaryKey": false, + "notNull": false + }, + "type": { + "name": "type", + "type": "machine_type", + "primaryKey": false, + "notNull": true, + "default": "'classic'" + } + }, + "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 + }, + "origin": { + "name": "origin", + "type": "workflow_run_origin", + "primaryKey": false, + "notNull": true, + "default": "'api'" + }, + "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" + } + }, + "machine_type": { + "name": "machine_type", + "values": { + "classic": "classic", + "runpod-serverless": "runpod-serverless" + } + }, + "workflow_run_origin": { + "name": "workflow_run_origin", + "values": { + "manual": "manual", + "api": "api" + } + }, + "workflow_run_status": { + "name": "workflow_run_status", + "values": { + "not-started": "not-started", + "running": "running", + "uploading": "uploading", + "success": "success", + "failed": "failed" + } + } + }, + "schemas": { + "comfyui_deploy": "comfyui_deploy" + }, + "_meta": { + "schemas": {}, + "tables": {}, + "columns": {} + } +} \ No newline at end of file diff --git a/web/drizzle/meta/_journal.json b/web/drizzle/meta/_journal.json index ea56dc6..85e9471 100644 --- a/web/drizzle/meta/_journal.json +++ b/web/drizzle/meta/_journal.json @@ -99,6 +99,20 @@ "when": 1703338425429, "tag": "0013_stormy_starbolt", "breakpoints": true + }, + { + "idx": 14, + "version": "5", + "when": 1703403388113, + "tag": "0014_short_sunfire", + "breakpoints": true + }, + { + "idx": 15, + "version": "5", + "when": 1703409502387, + "tag": "0015_simple_killmonger", + "breakpoints": true } ] } \ No newline at end of file diff --git a/web/package.json b/web/package.json index 4785091..22d78c6 100644 --- a/web/package.json +++ b/web/package.json @@ -42,6 +42,7 @@ "date-fns": "^3.0.5", "dayjs": "^1.11.10", "drizzle-orm": "^0.29.1", + "drizzle-zod": "^0.5.1", "jsonwebtoken": "^9.0.2", "lucide-react": "^0.294.0", "nanoid": "^5.0.4", diff --git a/web/src/components/InsertModal.tsx b/web/src/components/InsertModal.tsx new file mode 100644 index 0000000..cb6251e --- /dev/null +++ b/web/src/components/InsertModal.tsx @@ -0,0 +1,123 @@ +"use client"; + +import { LoadingIcon } from "./LoadingIcon"; +import { callServerPromise } from "@/components/callServerPromise"; +import AutoForm, { AutoFormSubmit } from "@/components/ui/auto-form"; +import type { FieldConfig } from "@/components/ui/auto-form/types"; +import { Button } from "@/components/ui/button"; +import { + Dialog, + DialogContent, + DialogDescription, + DialogHeader, + DialogTitle, + DialogTrigger, +} from "@/components/ui/dialog"; +import * as React from "react"; +import type { UnknownKeysParam, ZodObject, ZodRawShape, z } from "zod"; + +export function InsertModal< + K extends ZodRawShape, + Y extends UnknownKeysParam, + Z extends ZodObject +>(props: { + title: string; + description: string; + serverAction: (data: z.infer) => Promise; + formSchema: Z; + fieldConfig?: FieldConfig>; +}) { + const [open, setOpen] = React.useState(false); + const [isLoading, setIsLoading] = React.useState(false); + + return ( + + + + + + + {props.title} + {props.description} + + { + setIsLoading(true); + await callServerPromise(props.serverAction(data)); + setIsLoading(false); + setOpen(false); + }} + > +
+ + Save Changes + {isLoading && } + +
+ + + + ); +} + +export function UpdateModal< + K extends ZodRawShape, + Y extends UnknownKeysParam, + Z extends ZodObject +>(props: { + open: boolean; + setOpen: (open: boolean) => void; + title: string; + description: string; + data: z.infer; + serverAction: ( + data: z.infer & { + id: string; + } + ) => Promise; + formSchema: Z; + fieldConfig?: FieldConfig>; +}) { + // const [open, setOpen] = React.useState(false); + const [isLoading, setIsLoading] = React.useState(false); + + return ( + + {/* + {props.title} + */} + + + {props.title} + {props.description} + + { + setIsLoading(true); + await callServerPromise( + props.serverAction({ + ...data, + id: props.data.id, + }) + ); + setIsLoading(false); + props.setOpen(false); + }} + > +
+ + Save Changes + {isLoading && } + +
+ + + + ); +} diff --git a/web/src/components/LiveStatus.tsx b/web/src/components/LiveStatus.tsx index e3f74ed..8df690c 100644 --- a/web/src/components/LiveStatus.tsx +++ b/web/src/components/LiveStatus.tsx @@ -5,7 +5,7 @@ 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"; +import { useEffect } from "react"; export function LiveStatus({ run, @@ -16,23 +16,30 @@ 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; // const [view, setView] = useState(); - if (data?.json.event == "executing" && data.json.data.node == undefined) { - status = "success"; - } else if (data?.json.event == "executing") { + // if (data?.json.event == "executing" && data.json.data.node == undefined) { + // status = "success"; + // } else + if (data?.json.event == "executing") { status = "running"; + } else if (data?.json.event == "uploading") { + status = "uploading"; + } else if (data?.json.event == "success") { + status = "success"; + } else if (data?.json.event == "failed") { + status = "failed"; } const router = useRouter(); useEffect(() => { if (data?.json.event === "outputs_uploaded") { - router.refresh() + router.refresh(); } }, [data?.json.event]); diff --git a/web/src/components/MachineList.tsx b/web/src/components/MachineList.tsx index 8cc370c..fdca613 100644 --- a/web/src/components/MachineList.tsx +++ b/web/src/components/MachineList.tsx @@ -1,28 +1,11 @@ "use client"; import { getRelativeTime } from "../lib/getRelativeTime"; -import { LoadingIcon } from "./LoadingIcon"; +import { InsertModal, UpdateModal } from "./InsertModal"; import { callServerPromise } from "./callServerPromise"; -import { - FormControl, - FormField, - FormItem, - FormLabel, - FormMessage, - Form, -} from "./ui/form"; import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; import { Checkbox } from "@/components/ui/checkbox"; -import { - Dialog, - DialogContent, - DialogDescription, - DialogFooter, - DialogHeader, - DialogTitle, - DialogTrigger, -} from "@/components/ui/dialog"; import { DropdownMenu, DropdownMenuContent, @@ -39,14 +22,15 @@ import { TableHeader, TableRow, } from "@/components/ui/table"; -import type { MachineType } from "@/db/schema"; +import { type MachineType } from "@/db/schema"; +import { addMachineSchema } from "@/server/addMachineSchema"; import { addMachine, deleteMachine, disableMachine, enableMachine, + updateMachine, } from "@/server/curdMachine"; -import { zodResolver } from "@hookform/resolvers/zod"; import type { ColumnDef, ColumnFiltersState, @@ -63,8 +47,7 @@ import { } from "@tanstack/react-table"; import { ArrowUpDown, MoreHorizontal } from "lucide-react"; import * as React from "react"; -import { useForm } from "react-hook-form"; -import { z } from "zod"; +import { useState } from "react"; export type Machine = MachineType; @@ -127,6 +110,13 @@ export const columns: ColumnDef[] = [ ); }, }, + { + accessorKey: "type", + header: () =>
Type
, + cell: ({ row }) => { + return
{row.original.type}
; + }, + }, { accessorKey: "date", sortingFn: "datetime", @@ -154,6 +144,7 @@ export const columns: ColumnDef[] = [ enableHiding: false, cell: ({ row }) => { const machine = row.original; + const [open, setOpen] = useState(false); return ( @@ -191,10 +182,33 @@ export const columns: ColumnDef[] = [ Disable Machine )} - {/* - View customer - View payment details */} + setOpen(true)}> + Edit + + ); }, @@ -241,33 +255,12 @@ export function MachineList({ data }: { data: Machine[] }) { className="max-w-sm" />
- - {/* - - - - - {table - .getAllColumns() - .filter((column) => column.getCanHide()) - .map((column) => { - return ( - - column.toggleVisibility(!!value) - } - > - {column.id} - - ); - })} - - */} +
@@ -347,95 +340,3 @@ export function MachineList({ data }: { data: Machine[] }) {
); } - -const formSchema = z.object({ - name: z.string().min(1), - endpoint: z.string().min(1), -}); - -function AddMachinesDialog() { - const [open, setOpen] = React.useState(false); - const form = useForm>({ - resolver: zodResolver(formSchema), - defaultValues: { - name: "My Local Machine", - endpoint: "http://127.0.0.1:8188", - }, - }); - - return ( - - - - - -
- { - await addMachine(data.name, data.endpoint); - // await new Promise(resolve => setTimeout(resolve, 3000)); - setOpen(false); - })} - > - - Add Machines - - Add Comfyui machines to your account. - - -
- {/*
*/} - ( - - Name - - - - {/* - This is your public display name. - */} - - - )} - /> - - ( - - Endpoint - - - - {/* - This is your public display name. - */} - - - )} - /> -
- - - - - - -
- ); -} - -function AddWorkflowButton({ pending }: { pending: boolean }) { - // const { pending } = useFormStatus(); - return ( - - ); -} diff --git a/web/src/components/MachinesWS.tsx b/web/src/components/MachinesWS.tsx index 3aaff34..c923040 100644 --- a/web/src/components/MachinesWS.tsx +++ b/web/src/components/MachinesWS.tsx @@ -62,9 +62,11 @@ export function MachinesWSMain(props: {
Machine Status
- {props.machines.map((x) => ( - - ))} + {props.machines + .filter((x) => x.type === "classic") + .map((x) => ( + + ))}
); @@ -112,13 +114,14 @@ function MachineWS({ if (!lastMessage?.data) return; const message = JSON.parse(lastMessage.data); - console.log(message.event, message); + // console.log(message.event, message); if (message.data.sid) { setSid(message.data.sid); } if (message.data?.prompt_id) { + console.log(message.event, message); addData(message.data.prompt_id, message); } diff --git a/web/src/components/ui/auto-form/fields/object.tsx b/web/src/components/ui/auto-form/fields/object.tsx index 9e00789..d40d90e 100644 --- a/web/src/components/ui/auto-form/fields/object.tsx +++ b/web/src/components/ui/auto-form/fields/object.tsx @@ -1,28 +1,28 @@ -import * as z from "zod"; -import { useForm } from "react-hook-form"; -import { FieldConfig, FieldConfigItem } from "../types"; import { Accordion, AccordionContent, AccordionItem, AccordionTrigger, } from "../../accordion"; +import { FormField } from "../../form"; +import { DEFAULT_ZOD_HANDLERS, INPUT_COMPONENTS } from "../config"; +import type { FieldConfig, FieldConfigItem } from "../types"; import { beautifyObjectName, getBaseSchema, getBaseType, zodToHtmlInputProps, } from "../utils"; -import { FormField } from "../../form"; -import { DEFAULT_ZOD_HANDLERS, INPUT_COMPONENTS } from "../config"; import AutoFormArray from "./array"; +import type { useForm } from "react-hook-form"; +import type * as z from "zod"; function DefaultParent({ children }: { children: React.ReactNode }) { return <>{children}; } export default function AutoFormObject< - SchemaType extends z.ZodObject, + SchemaType extends z.ZodObject >({ schema, form, diff --git a/web/src/components/ui/auto-form/index.tsx b/web/src/components/ui/auto-form/index.tsx index 668cd16..fa5d7e3 100644 --- a/web/src/components/ui/auto-form/index.tsx +++ b/web/src/components/ui/auto-form/index.tsx @@ -1,20 +1,16 @@ "use client"; -import React from "react"; -import { z } from "zod"; -import { Form } from "../form"; -import { DefaultValues, useForm } from "react-hook-form"; -import { zodResolver } from "@hookform/resolvers/zod"; import { Button } from "../button"; +import { Form } from "../form"; +import type { FieldConfig } from "./types"; +import type { ZodObjectOrWrapped } from "./utils"; +import { getDefaultValues, getObjectFormSchema } from "./utils"; +import AutoFormObject from "@/components/ui/auto-form/fields/object"; import { cn } from "@/lib/utils"; - -import { FieldConfig } from "./types"; -import { - ZodObjectOrWrapped, - getDefaultValues, - getObjectFormSchema, -} from "./utils"; -import AutoFormObject from "./fields/object"; +import { zodResolver } from "@hookform/resolvers/zod"; +import type { DefaultValues } from "react-hook-form"; +import { useForm } from "react-hook-form"; +import type { z } from "zod"; export function AutoFormSubmit({ children }: { children?: React.ReactNode }) { return ; diff --git a/web/src/db/schema.ts b/web/src/db/schema.ts index a895442..d237818 100644 --- a/web/src/db/schema.ts +++ b/web/src/db/schema.ts @@ -9,6 +9,7 @@ import { pgEnum, boolean, } from "drizzle-orm/pg-core"; +import { createInsertSchema } from "drizzle-zod"; import { z } from "zod"; export const dbSchema = pgSchema("comfyui_deploy"); @@ -100,6 +101,11 @@ export const workflowRunOrigin = pgEnum("workflow_run_origin", [ "api", ]); +export const machinesType = pgEnum("machine_type", [ + "classic", + "runpod-serverless", +]); + // We still want to keep the workflow run record. export const workflowRunsTable = dbSchema.table("workflow_runs", { id: uuid("id").primaryKey().defaultRandom().notNull(), @@ -178,6 +184,14 @@ export const machinesTable = dbSchema.table("machines", { created_at: timestamp("created_at").defaultNow().notNull(), updated_at: timestamp("updated_at").defaultNow().notNull(), disabled: boolean("disabled").default(false).notNull(), + auth_token: text("auth_token"), + type: machinesType("type").notNull().default("classic"), +}); + +export const insertMachineSchema = createInsertSchema(machinesTable, { + name: (schema) => schema.name.default("My Machine"), + endpoint: (schema) => schema.endpoint.default("http://127.0.0.1:8188"), + type: (schema) => schema.type.default("classic"), }); export const deploymentsTable = dbSchema.table("deployments", { diff --git a/web/src/server/addMachineSchema.ts b/web/src/server/addMachineSchema.ts new file mode 100644 index 0000000..1119930 --- /dev/null +++ b/web/src/server/addMachineSchema.ts @@ -0,0 +1,8 @@ +import { insertMachineSchema } from "@/db/schema"; + +export const addMachineSchema = insertMachineSchema.pick({ + name: true, + endpoint: true, + type: true, + auth_token: true, +}); diff --git a/web/src/server/createRun.ts b/web/src/server/createRun.ts index d0c56ec..7ba530a 100644 --- a/web/src/server/createRun.ts +++ b/web/src/server/createRun.ts @@ -7,6 +7,7 @@ import { ComfyAPI_Run } from "@/types/ComfyAPI_Run"; import { and, eq } from "drizzle-orm"; import { revalidatePath } from "next/cache"; import "server-only"; +import { v4 } from "uuid"; export const createRun = withServerPromise( async ( @@ -37,8 +38,6 @@ export const createRun = withServerPromise( throw new Error("Workflow version not found"); } - const comfyui_endpoint = `${machine.endpoint}/comfyui-deploy/run`; - const workflow_api = workflow_version_data.workflow_api; // Replace the inputs @@ -52,33 +51,57 @@ export const createRun = withServerPromise( } } - const body = { + let prompt_id: string | undefined = undefined; + const shareData = { 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); - // Sending to comfyui - const _result = await fetch(comfyui_endpoint, { - method: "POST", - body: bodyJson, - cache: "no-store", - }); - - if (!_result.ok) { - throw new Error(`Error creating run, ${_result.statusText}`); + switch (machine.type) { + case "runpod-serverless": + prompt_id = v4(); + const data = { + input: { + ...shareData, + prompt_id: prompt_id, + }, + }; + console.log(data); + const __result = await fetch(`${machine.endpoint}/run`, { + method: "POST", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${machine.auth_token}`, + }, + body: JSON.stringify(data), + cache: "no-store", + }); + console.log(__result); + if (!__result.ok) + throw new Error(`Error creating run, ${__result.statusText}`); + console.log(data, __result); + break; + case "classic": + const body = shareData; + const comfyui_endpoint = `${machine.endpoint}/comfyui-deploy/run`; + const _result = await fetch(comfyui_endpoint, { + method: "POST", + body: JSON.stringify(body), + cache: "no-store", + }); + if (!_result.ok) + throw new Error(`Error creating run, ${_result.statusText}`); + const result = await ComfyAPI_Run.parseAsync(await _result.json()); + prompt_id = result.prompt_id; + break; } - const result = await ComfyAPI_Run.parseAsync(await _result.json()); - // Add to our db const workflow_run = await db .insert(workflowRunsTable) .values({ - id: result.prompt_id, + id: prompt_id, workflow_id: workflow_version_data.workflow_id, workflow_version_id: workflow_version_data.id, workflow_inputs: inputs, diff --git a/web/src/server/curdMachine.ts b/web/src/server/curdMachine.ts index 45feab4..11a89ea 100644 --- a/web/src/server/curdMachine.ts +++ b/web/src/server/curdMachine.ts @@ -1,5 +1,6 @@ "use server"; +import type { addMachineSchema } from "./addMachineSchema"; import { withServerPromise } from "./withServerPromise"; import { db } from "@/db/db"; import { machinesTable } from "@/db/schema"; @@ -7,6 +8,7 @@ import { auth } from "@clerk/nextjs"; import { and, eq } from "drizzle-orm"; import { revalidatePath } from "next/cache"; import "server-only"; +import type { z } from "zod"; export async function getMachines() { const { userId } = auth(); @@ -20,17 +22,36 @@ export async function getMachines() { return machines; } -export async function addMachine(name: string, endpoint: string) { - const { userId } = auth(); - if (!userId) throw new Error("No user id"); - console.log(name, endpoint); - await db.insert(machinesTable).values({ - user_id: userId, - name, - endpoint, - }); - revalidatePath("/machines"); -} +export const addMachine = withServerPromise( + async ({ name, endpoint, type }: z.infer) => { + const { userId } = auth(); + if (!userId) return { error: "No user id" }; + console.log(name, endpoint); + await db.insert(machinesTable).values({ + user_id: userId, + name, + endpoint, + type, + }); + revalidatePath("/machines"); + return { message: "Machine Added" }; + } +); + +export const updateMachine = withServerPromise( + async ({ + id, + ...data + }: z.infer & { + id: string; + }) => { + const { userId } = auth(); + if (!userId) return { error: "No user id" }; + await db.update(machinesTable).set(data).where(eq(machinesTable.id, id)); + revalidatePath("/machines"); + return { message: "Machine Updated" }; + } +); export const deleteMachine = withServerPromise( async (machine_id: string): Promise<{ message: string }> => {