diff --git a/web/src/app/api/upload/route.ts b/web/src/app/api/upload/route.ts
index 5ec0065..139d48d 100644
--- a/web/src/app/api/upload/route.ts
+++ b/web/src/app/api/upload/route.ts
@@ -1,6 +1,11 @@
-import { parseDataSafe } from "../../../lib/parseDataSafe";
 import { db } from "@/db/db";
-import { workflowTable, workflowVersionTable } from "@/db/schema";
+import {
+  workflowAPIType,
+  workflowTable,
+  workflowType,
+  workflowVersionTable,
+} from "@/db/schema";
+import { parseDataSafe } from "@/lib/parseDataSafe";
 import { sql } from "drizzle-orm";
 import { NextResponse } from "next/server";
 import { z } from "zod";
@@ -15,8 +20,8 @@ const UploadRequest = z.object({
   user_id: z.string(),
   workflow_id: z.string().optional(),
   workflow_name: z.string().optional(),
-  workflow: z.any(),
-  workflow_api: z.any(),
+  workflow: workflowType,
+  workflow_api: workflowAPIType,
 });
 
 export async function OPTIONS(request: Request) {
@@ -41,7 +46,15 @@ export async function POST(request: Request) {
 
   if (!data || error) return error;
 
-  let { user_id, workflow, workflow_api, workflow_id, workflow_name } = data;
+  const {
+    user_id,
+    workflow,
+    workflow_api,
+    workflow_id: _workflow_id,
+    workflow_name,
+  } = data;
+
+  let workflow_id = _workflow_id;
 
   let version = -1;
 
@@ -49,7 +62,7 @@ export async function POST(request: Request) {
   try {
     if ((!workflow_id || workflow_id.length == 0) && workflow_name) {
       // Create a new parent workflow
-      const workflow = await db
+      const workflow_parent = await db
         .insert(workflowTable)
         .values({
           user_id,
@@ -57,7 +70,7 @@ export async function POST(request: Request) {
         })
         .returning();
 
-      workflow_id = workflow[0].id;
+      workflow_id = workflow_parent[0].id;
 
       // Create a new version
       const data = await db
@@ -76,7 +89,7 @@ export async function POST(request: Request) {
         .insert(workflowVersionTable)
         .values({
           workflow_id,
-          workflow,
+          workflow: workflow,
           workflow_api,
           // version: sql`${workflowVersionTable.version} + 1`,
           version: sql`(
diff --git a/web/src/db/schema.ts b/web/src/db/schema.ts
index 8d814e6..1d999be 100644
--- a/web/src/db/schema.ts
+++ b/web/src/db/schema.ts
@@ -8,6 +8,7 @@ import {
   jsonb,
   pgEnum,
 } from "drizzle-orm/pg-core";
+import { z } from "zod";
 
 export const dbSchema = pgSchema("comfyui_deploy");
 
@@ -35,21 +36,25 @@ export const workflowRelations = relations(workflowTable, ({ many }) => ({
   versions: many(workflowVersionTable),
 }));
 
-export type WorkflowJSONType = {
-  nodes: {
-    id: string;
-    type: string;
-    widgets_values: any[];
-  }[];
-};
+export const workflowType = z.any();
+// export const workflowType = z.object({
+//   last_node_id: z.number(),
+//   last_link_id: z.number(),
+//   nodes: z.array(
+//     z.object({
+//       id: z.number(),
+//       type: z.string(),
+//       widgets_values: z.array(z.any()),
+//     })
+//   ),
+// });
 
-export type WorkflowAPIType = Record<
-  string,
-  {
-    inputs: Record<string, string>;
-    class_type: string;
-  }
->;
+export const workflowAPIType = z.record(
+  z.object({
+    inputs: z.record(z.any()),
+    class_type: z.string(),
+  })
+);
 
 export const workflowVersionTable = dbSchema.table("workflow_versions", {
   workflow_id: uuid("workflow_id")
@@ -58,8 +63,8 @@ export const workflowVersionTable = dbSchema.table("workflow_versions", {
       onDelete: "cascade",
     }),
   id: uuid("id").primaryKey().defaultRandom().notNull(),
-  workflow: jsonb("workflow").$type<WorkflowJSONType>(),
-  workflow_api: jsonb("workflow_api").$type<WorkflowAPIType>(),
+  workflow: jsonb("workflow").$type<z.infer<typeof workflowType>>(),
+  workflow_api: jsonb("workflow_api").$type<z.infer<typeof workflowAPIType>>(),
   version: integer("version").notNull(),
 
   created_at: timestamp("created_at").defaultNow().notNull(),