fix: handle file upload status correctly, add serverless machine type

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}")
 

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;

web/drizzle/0015_simple_killmonger.sql

@@ -0,0 +1 @@
+ALTER TABLE "comfyui_deploy"."machines" ADD COLUMN "auth_token" text;

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": {}
+  }
+}

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": {}
+  }
+}

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
     }
   ]
 }

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",

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<K, Y>
+>(props: {
+  title: string;
+  description: string;
+  serverAction: (data: z.infer<Z>) => Promise<unknown>;
+  formSchema: Z;
+  fieldConfig?: FieldConfig<z.infer<Z>>;
+}) {
+  const [open, setOpen] = React.useState(false);
+  const [isLoading, setIsLoading] = React.useState(false);
+
+  return (
+    <Dialog open={open} onOpenChange={setOpen}>
+      <DialogTrigger asChild>
+        <Button variant="default" className="">
+          {props.title}
+        </Button>
+      </DialogTrigger>
+      <DialogContent className="sm:max-w-[425px]">
+        <DialogHeader>
+          <DialogTitle>{props.title}</DialogTitle>
+          <DialogDescription>{props.description}</DialogDescription>
+        </DialogHeader>
+        <AutoForm
+          fieldConfig={props.fieldConfig}
+          formSchema={props.formSchema}
+          onSubmit={async (data) => {
+            setIsLoading(true);
+            await callServerPromise(props.serverAction(data));
+            setIsLoading(false);
+            setOpen(false);
+          }}
+        >
+          <div className="flex justify-end">
+            <AutoFormSubmit>
+              Save Changes
+              <span className="ml-2">{isLoading && <LoadingIcon />}</span>
+            </AutoFormSubmit>
+          </div>
+        </AutoForm>
+      </DialogContent>
+    </Dialog>
+  );
+}
+
+export function UpdateModal<
+  K extends ZodRawShape,
+  Y extends UnknownKeysParam,
+  Z extends ZodObject<K, Y>
+>(props: {
+  open: boolean;
+  setOpen: (open: boolean) => void;
+  title: string;
+  description: string;
+  data: z.infer<Z>;
+  serverAction: (
+    data: z.infer<Z> & {
+      id: string;
+    }
+  ) => Promise<unknown>;
+  formSchema: Z;
+  fieldConfig?: FieldConfig<z.infer<Z>>;
+}) {
+  // const [open, setOpen] = React.useState(false);
+  const [isLoading, setIsLoading] = React.useState(false);
+
+  return (
+    <Dialog open={props.open} onOpenChange={props.setOpen}>
+      {/* <DialogTrigger asChild>
+        <DropdownMenuItem>{props.title}</DropdownMenuItem>
+      </DialogTrigger> */}
+      <DialogContent className="sm:max-w-[425px]">
+        <DialogHeader>
+          <DialogTitle>{props.title}</DialogTitle>
+          <DialogDescription>{props.description}</DialogDescription>
+        </DialogHeader>
+        <AutoForm
+          fieldConfig={props.fieldConfig}
+          formSchema={props.formSchema}
+          onSubmit={async (data) => {
+            setIsLoading(true);
+            await callServerPromise(
+              props.serverAction({
+                ...data,
+                id: props.data.id,
+              })
+            );
+            setIsLoading(false);
+            props.setOpen(false);
+          }}
+        >
+          <div className="flex justify-end">
+            <AutoFormSubmit>
+              Save Changes
+              <span className="ml-2">{isLoading && <LoadingIcon />}</span>
+            </AutoFormSubmit>
+          </div>
+        </AutoForm>
+      </DialogContent>
+    </Dialog>
+  );
+}

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<any>();
-  if (data?.json.event == "executing" && data.json.data.node == undefined) {
+  // if (data?.json.event == "executing" && data.json.data.node == undefined) {
-    status = "success";
+  //   status = "success";
-  } else if (data?.json.event == "executing") {
+  // } 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]);
 

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 { useState } from "react";
-import { z } from "zod";
 
 export type Machine = MachineType;
 
@@ -127,6 +110,13 @@ export const columns: ColumnDef<Machine>[] = [
       );
     },
   },
+  {
+    accessorKey: "type",
+    header: () => <div className="text-left">Type</div>,
+    cell: ({ row }) => {
+      return <div className="text-left font-medium">{row.original.type}</div>;
+    },
+  },
   {
     accessorKey: "date",
     sortingFn: "datetime",
@@ -154,6 +144,7 @@ export const columns: ColumnDef<Machine>[] = [
     enableHiding: false,
     cell: ({ row }) => {
       const machine = row.original;
+      const [open, setOpen] = useState(false);
 
       return (
         <DropdownMenu>
@@ -191,10 +182,33 @@ export const columns: ColumnDef<Machine>[] = [
                 Disable Machine
               </DropdownMenuItem>
             )}
-            {/* <DropdownMenuSeparator />
+            <DropdownMenuItem onClick={() => setOpen(true)}>
-            <DropdownMenuItem>View customer</DropdownMenuItem>
+              Edit
-            <DropdownMenuItem>View payment details</DropdownMenuItem> */}
+            </DropdownMenuItem>
           </DropdownMenuContent>
+          <UpdateModal
+            fieldConfig={{
+              name: {
+                inputProps: { defaultValue: machine.name },
+              },
+              endpoint: {
+                inputProps: { defaultValue: machine.endpoint },
+              },
+              type: {
+                inputProps: { defaultValue: machine.type },
+              },
+              auth_token: {
+                inputProps: { defaultValue: machine.auth_token ?? "" },
+              },
+            }}
+            data={machine}
+            open={open}
+            setOpen={setOpen}
+            title="Edit"
+            description="Edit machines"
+            serverAction={updateMachine}
+            formSchema={addMachineSchema}
+          />
         </DropdownMenu>
       );
     },
@@ -241,33 +255,12 @@ export function MachineList({ data }: { data: Machine[] }) {
           className="max-w-sm"
         />
         <div className="ml-auto flex gap-2">
-          <AddMachinesDialog />
+          <InsertModal
-          {/* <DropdownMenu>
+            title="Add Machine"
-            <DropdownMenuTrigger asChild>
+            description="Add Comfyui machines to your account."
-              <Button variant="outline" className="">
+            serverAction={addMachine}
-                Columns <ChevronDown className="ml-2 h-4 w-4" />
+            formSchema={addMachineSchema}
-              </Button>
+          />
-            </DropdownMenuTrigger>
-            <DropdownMenuContent align="end">
-              {table
-                .getAllColumns()
-                .filter((column) => column.getCanHide())
-                .map((column) => {
-                  return (
-                    <DropdownMenuCheckboxItem
-                      key={column.id}
-                      className="capitalize"
-                      checked={column.getIsVisible()}
-                      onCheckedChange={(value) =>
-                        column.toggleVisibility(!!value)
-                      }
-                    >
-                      {column.id}
-                    </DropdownMenuCheckboxItem>
-                  );
-                })}
-            </DropdownMenuContent>
-          </DropdownMenu> */}
         </div>
       </div>
       <div className="rounded-md border overflow-x-auto w-full">
@@ -347,95 +340,3 @@ export function MachineList({ data }: { data: Machine[] }) {
     </div>
   );
 }
-
-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<z.infer<typeof formSchema>>({
-    resolver: zodResolver(formSchema),
-    defaultValues: {
-      name: "My Local Machine",
-      endpoint: "http://127.0.0.1:8188",
-    },
-  });
-
-  return (
-    <Dialog open={open} onOpenChange={setOpen}>
-      <DialogTrigger asChild>
-        <Button variant="default" className="">
-          Add Machines
-        </Button>
-      </DialogTrigger>
-      <DialogContent className="sm:max-w-[425px]">
-        <Form {...form}>
-          <form
-            onSubmit={form.handleSubmit(async (data) => {
-              await addMachine(data.name, data.endpoint);
-              // await new Promise(resolve => setTimeout(resolve, 3000));
-              setOpen(false);
-            })}
-          >
-            <DialogHeader>
-              <DialogTitle>Add Machines</DialogTitle>
-              <DialogDescription>
-                Add Comfyui machines to your account.
-              </DialogDescription>
-            </DialogHeader>
-            <div className="grid gap-4 py-4">
-              {/* <div className="grid grid-cols-4 items-center gap-4"> */}
-              <FormField
-                control={form.control}
-                name="name"
-                render={({ field }) => (
-                  <FormItem>
-                    <FormLabel>Name</FormLabel>
-                    <FormControl>
-                      <Input {...field} />
-                    </FormControl>
-                    {/* <FormDescription>
-                      This is your public display name.
-                    </FormDescription> */}
-                    <FormMessage />
-                  </FormItem>
-                )}
-              />
-
-              <FormField
-                control={form.control}
-                name="endpoint"
-                render={({ field }) => (
-                  <FormItem>
-                    <FormLabel>Endpoint</FormLabel>
-                    <FormControl>
-                      <Input {...field} />
-                    </FormControl>
-                    {/* <FormDescription>
-                      This is your public display name.
-                    </FormDescription> */}
-                    <FormMessage />
-                  </FormItem>
-                )}
-              />
-            </div>
-            <DialogFooter>
-              <AddWorkflowButton pending={form.formState.isSubmitting} />
-            </DialogFooter>
-          </form>
-        </Form>
-      </DialogContent>
-    </Dialog>
-  );
-}
-
-function AddWorkflowButton({ pending }: { pending: boolean }) {
-  // const { pending } = useFormStatus();
-  return (
-    <Button type="submit" disabled={pending}>
-      Save changes {pending && <LoadingIcon />}
-    </Button>
-  );
-}

web/src/components/MachinesWS.tsx

@@ -62,9 +62,11 @@ export function MachinesWSMain(props: {
     <div className="flex flex-col gap-2 mt-4">
       Machine Status
       <div className="flex flex-wrap gap-2">
-        {props.machines.map((x) => (
+        {props.machines
-          <MachineWS key={x.id} machine={x} />
+          .filter((x) => x.type === "classic")
-        ))}
+          .map((x) => (
+            <MachineWS key={x.id} machine={x} />
+          ))}
       </div>
     </div>
   );
@@ -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);
     }
 

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<any, any>,
+  SchemaType extends z.ZodObject<any, any>
 >({
   schema,
   form,

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 { zodResolver } from "@hookform/resolvers/zod";
-import { FieldConfig } from "./types";
+import type { DefaultValues } from "react-hook-form";
-import {
+import { useForm } from "react-hook-form";
-  ZodObjectOrWrapped,
+import type { z } from "zod";
-  getDefaultValues,
-  getObjectFormSchema,
-} from "./utils";
-import AutoFormObject from "./fields/object";
 
 export function AutoFormSubmit({ children }: { children?: React.ReactNode }) {
   return <Button type="submit">{children ?? "Submit"}</Button>;

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", {

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,
+});

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
+    switch (machine.type) {
-    const _result = await fetch(comfyui_endpoint, {
+      case "runpod-serverless":
-      method: "POST",
+        prompt_id = v4();
-      body: bodyJson,
+        const data = {
-      cache: "no-store",
+          input: {
-    });
+            ...shareData,
-
+            prompt_id: prompt_id,
-    if (!_result.ok) {
+          },
-      throw new Error(`Error creating run, ${_result.statusText}`);
+        };
+        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,

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) {
+export const addMachine = withServerPromise(
-  const { userId } = auth();
+  async ({ name, endpoint, type }: z.infer<typeof addMachineSchema>) => {
-  if (!userId) throw new Error("No user id");
+    const { userId } = auth();
-  console.log(name, endpoint);
+    if (!userId) return { error: "No user id" };
-  await db.insert(machinesTable).values({
+    console.log(name, endpoint);
-    user_id: userId,
+    await db.insert(machinesTable).values({
-    name,
+      user_id: userId,
-    endpoint,
+      name,
-  });
+      endpoint,
-  revalidatePath("/machines");
+      type,
-}
+    });
+    revalidatePath("/machines");
+    return { message: "Machine Added" };
+  }
+);
+
+export const updateMachine = withServerPromise(
+  async ({
+    id,
+    ...data
+  }: z.infer<typeof addMachineSchema> & {
+    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 }> => {