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

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) {
-    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]);
 

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<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>View customer</DropdownMenuItem>
-            <DropdownMenuItem>View payment details</DropdownMenuItem> */}
+            <DropdownMenuItem onClick={() => setOpen(true)}>
+              Edit
+            </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 />
-          {/* <DropdownMenu>
-            <DropdownMenuTrigger asChild>
-              <Button variant="outline" className="">
-                Columns <ChevronDown className="ml-2 h-4 w-4" />
-              </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> */}
+          <InsertModal
+            title="Add Machine"
+            description="Add Comfyui machines to your account."
+            serverAction={addMachine}
+            formSchema={addMachineSchema}
+          />
         </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) => (
-          <MachineWS key={x.id} machine={x} />
-        ))}
+        {props.machines
+          .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 { 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 <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
-    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,

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<typeof addMachineSchema>) => {
+    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<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 }> => {