From eac8d709f9f756b202d7ac837761833feb5624db Mon Sep 17 00:00:00 2001
From: BennyKok <itechbenny@gmail.com>
Date: Fri, 15 Dec 2023 12:03:40 +0800
Subject: [PATCH] fix: run with external inputs

---
 web/src/db/schema.ts        | 20 ++++++++++++++++++--
 web/src/server/createRun.ts | 11 +++++++----
 2 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/web/src/db/schema.ts b/web/src/db/schema.ts
index 9883e13..e7ac3de 100644
--- a/web/src/db/schema.ts
+++ b/web/src/db/schema.ts
@@ -35,6 +35,22 @@ export const workflowRelations = relations(workflowTable, ({ many }) => ({
   versions: many(workflowVersionTable),
 }));
 
+export type WorkflowJSONType = {
+  nodes: {
+    id: string;
+    type: string;
+    widgets_values: any[];
+  }[];
+};
+
+export type WorkflowAPIType = Record<
+  string,
+  {
+    inputs: Record<string, string>;
+    class_type: string;
+  }
+>;
+
 export const workflowVersionTable = dbSchema.table("workflow_versions", {
   workflow_id: uuid("workflow_id")
     .notNull()
@@ -42,8 +58,8 @@ export const workflowVersionTable = dbSchema.table("workflow_versions", {
       onDelete: "cascade",
     }),
   id: uuid("id").primaryKey().defaultRandom().notNull(),
-  workflow: jsonb("workflow").$type<any>(),
-  workflow_api: jsonb("workflow_api").$type<any>(),
+  workflow: jsonb("workflow").$type<WorkflowJSONType>(),
+  workflow_api: jsonb("workflow_api").$type<WorkflowAPIType>(),
   version: integer("version").notNull(),
 
   created_at: timestamp("created_at").defaultNow().notNull(),
diff --git a/web/src/server/createRun.ts b/web/src/server/createRun.ts
index 0e3294e..9ee63d2 100644
--- a/web/src/server/createRun.ts
+++ b/web/src/server/createRun.ts
@@ -47,12 +47,16 @@ export async function createRun(
 
   const comfyui_endpoint = `${machine.endpoint}/comfyui-deploy/run`;
 
-  let workflow_api = workflow_version_data.workflow_api;
+  const workflow_api = workflow_version_data.workflow_api;
 
   // Replace the inputs
-  if (inputs) {
+  if (inputs && workflow_api) {
     for (const key in inputs) {
-      workflow_api = workflow_api.replace(`"${key}"`, `"${inputs[key]}"`);
+      Object.entries(workflow_api).forEach(([_, node]) => {
+        if (node.inputs["name"] === key) {
+          node.inputs["name"] = inputs[key];
+        }
+      });
     }
   }
 
@@ -66,7 +70,6 @@ export async function createRun(
       workflow_api: workflow_api,
       status_endpoint: `${origin}/api/update-run`,
       file_upload_endpoint: `${origin}/api/file-upload`,
-      inputs: inputs,
     }),
   }).then(async (res) => ComfyAPI_Run.parseAsync(await res.json()));
   // .catch((error) => {