fix: type

* fix: attempt fixing timeout

* be validation work

* arbitrary model input, BE validation, error record creation with error logs during potential failure points

* remove unused type

---------

Co-authored-by: bennykok <[email protected]>

builder/modal-builder/README.md

@@ -56,5 +56,5 @@ fly launch
 ```
 if not, run this instead
 ```
-fly deploy
+fly deploy -c "toml file"
 ```

builder/modal-builder/src/template/app.py

@@ -158,7 +158,7 @@ image = Image.debian_slim()
 
 target_image = image if deploy_test else dockerfile_image
 
-@stub.cls(image=target_image, gpu=config["gpu"] ,volumes=volumes, timeout=60 * 10, container_idle_timeout=60 * 5)
+@stub.cls(image=target_image, gpu=config["gpu"] ,volumes=volumes, timeout=60 * 10, container_idle_timeout=60)
 class ComfyDeployRunner:
 
     @enter()

builder/modal-builder/src/volume-builder/app.py

@@ -42,7 +42,7 @@ def download_model(volume_name, download_config):
 
     volume_base_path = vol_name_to_path[volume_name]
     model_store_path = os.path.join(volume_base_path, folder_path)
-    modified_download_url = download_url + ("&" if "?" in download_url else "?") + "token=" + civitai_key
+    modified_download_url = download_url + ("&" if "?" in download_url else "?") + "token=" + civitai_key # civitai requires auth
     print('downloading', modified_download_url)
 
     subprocess.run(["wget", modified_download_url , "--content-disposition", "-P", model_store_path])

web/drizzle/0046_complex_mentallo.sql

@@ -0,0 +1 @@
+ALTER TYPE "model_upload_type" ADD VALUE 'download_url';

web/drizzle/0047_gifted_starbolt.sql

@@ -0,0 +1 @@
+ALTER TYPE "model_upload_type" ADD VALUE 'download-url';

web/drizzle/meta/_journal.json

@@ -323,6 +323,20 @@
       "when": 1706336448134,
       "tag": "0045_careful_cerise",
       "breakpoints": true
+    },
+    {
+      "idx": 46,
+      "version": "5",
+      "when": 1706383154642,
+      "tag": "0046_complex_mentallo",
+      "breakpoints": true
+    },
+    {
+      "idx": 47,
+      "version": "5",
+      "when": 1706384528895,
+      "tag": "0047_gifted_starbolt",
+      "breakpoints": true
     }
   ]
 }

web/src/app/(app)/api/stripe/checkout/route.tsx

@@ -36,6 +36,7 @@ export async function GET(req: Request) {
       orgId: orgId ?? null,
       plan: plan,
     },
+    allow_promotion_codes: true,
     client_reference_id: orgId ?? userId,
     customer_email: user.emailAddresses[0].emailAddress,
     mode: "subscription",

web/src/app/(app)/machines/page.tsx

@@ -1,7 +1,9 @@
 import { AccessType } from "../../../lib/AccessType";
 import { MachineList } from "@/components/MachineList";
+import { SubscriptionProvider } from "@/components/useCurrentPlan";
 import { db } from "@/db/db";
 import { machinesTable } from "@/db/schema";
+import { getCurrentPlanWithAuth } from "@/server/getCurrentPlan";
 import { auth } from "@clerk/nextjs";
 import { clerkClient } from "@clerk/nextjs/server";
 import { desc, eq, isNull, and } from "drizzle-orm";
@@ -19,6 +21,8 @@ async function MachineListServer() {
 
   const user = await clerkClient.users.getUser(userId);
 
+  const sub = await getCurrentPlanWithAuth();
+
   const machines = await db.query.machinesTable.findMany({
     orderBy: desc(machinesTable.updated_at),
     where:
@@ -29,11 +33,12 @@ async function MachineListServer() {
 
   return (
     <div className="w-full">
-      {/* <div>Machines</div> */}
+      <SubscriptionProvider sub={sub}>
         <MachineList
           data={machines}
           userMetadata={AccessType.parse(user.privateMetadata ?? {})}
         />
+      </SubscriptionProvider>
     </div>
   );
 }

web/src/app/(docs)/docs/install/page.mdx

@@ -1,12 +1,12 @@
 export const metadata = {
   title: 'Quickstart',
   description:
-    'This guide will get you all set up and ready to use the Protocol API. We’ll cover how to get started an API client and how to make your first API request.',
+    'This guide will get you all set up and ready to use Comfy Deploy. We’ll cover how to get started an API client and how to make your first API request.',
 }
 
-# Getting stated
+# Getting started
 
-Install Comfy Deploy's plugin on your local machine to get started with deploying workflow.
+Install Comfy Deploy's plugin on your local ComfyUI to get started with deploying workflow.
 
 <CodeGroup>
 

web/src/app/(docs)/docs/videos/page.mdx

@@ -0,0 +1,58 @@
+export const metadata = {
+  title: 'Comfy Deploy Video tutorials',
+  description:
+    'Get started with API integration to run any deploy ComfyUI workflow',
+}
+
+**This is a collection of video tutorials on questions that have come up in our [discord channel](https://discord.gg/Zrar7yP6MJ).**
+
+# Comfy Deploy workflow walkthrough (comprehensive)
+
+Local workflow, to production API
+- Set inputs with Comfy Deploy `External Text` nodes
+- Deploy the workflow 
+- Create a machine
+- Deploy an API endpoint for the workflow
+- Test the workflow in Comfy Deploy
+
+<video controls>
+  <source src="https://pub-1a2bca0642c24fcfb84ce8d8415958d3.r2.dev/comfydeploy_base-sd-setup.mp4" type="video/mp4"/>
+  Your browser does not support the video tag.
+</video>
+
+
+
+# Test with Docs
+Shows how you can use the <a href='/docs/endpoints'>endpoints docs</a> to test out the API!
+
+<video controls>
+  <source src="https://pub-1a2bca0642c24fcfb84ce8d8415958d3.r2.dev/comfydeploy_interactive_docs.mp4" type="video/mp4"/>
+  Your browser does not support the video tag.
+</video>
+
+## Intermediate IPAdaptor example run 
+Shows off a semi-complex workflow working on Comfy Deploy
+
+<video controls>
+  <source src="https://pub-1a2bca0642c24fcfb84ce8d8415958d3.r2.dev/ipadapters-test.mp4" type="video/mp4"/>
+  Your browser does not support the video tag.
+</video>
+
+## Animate diff on Comfy Deploy
+Shows off animate diff working on Comfy Deploy
+
+<video controls>
+  <source src="https://pub-1a2bca0642c24fcfb84ce8d8415958d3.r2.dev/comfydeploy-animediff.mp4" type="video/mp4"/>
+  Your browser does not support the video tag.
+</video>
+
+## Cloning workflows
+What does cloning a workflow look like
+
+<video controls>
+  <source src="https://pub-1a2bca0642c24fcfb84ce8d8415958d3.r2.dev/comfydeploy-clone-template.mp4" type="video/mp4"/>
+  Your browser does not support the video tag.
+</video>
+
+## [Install custom nodes from any github repo (Loom)](https://www.loom.com/share/c2025c5060e348839ade6a3a9d96441d?sid=6a57f779-0209-4341-a5fd-f99769cdb162)
+- This is only neccessary if the custom node is not in the ComfyUI Manager node list

web/src/components/CurrentPlanContextType.tsx

@@ -0,0 +1,24 @@
+"use client";
+import { getCurrentPlanWithAuth } from "@/server/getCurrentPlan";
+import * as React from "react";
+import { createContext } from "react";
+
+export type CurrentPlanContextType = Awaited<
+  ReturnType<typeof getCurrentPlanWithAuth>
+>;
+export const CurrentPlanContext = createContext<
+  CurrentPlanContextType | undefined
+>(undefined);
+export function SubscriptionProvider({
+  sub,
+  children,
+}: {
+  sub: CurrentPlanContextType;
+  children: React.ReactNode;
+}) {
+  return (
+    <CurrentPlanContext.Provider value={sub}>
+      {children}
+    </CurrentPlanContext.Provider>
+  );
+}

web/src/components/DeploymentDisplay.tsx

@@ -118,15 +118,12 @@ export function DeploymentDisplay({
             </TabsList>
             <TabsContent className="flex flex-col gap-2 !mt-0" value="client">
               <div>
-                Copy and paste the ComfyDeployClient form&nbsp;
+                Install the node comfydeploy SDK
-                <a
+                <CodeBlock
-                  href="https://github.com/BennyKok/comfyui-deploy-next-example/blob/main/src/lib/comfy-deploy.ts"
+                  lang="bash"
-                  className="text-blue-500 hover:underline"
+                  code={`npm i comfydeploy`}
-                  target="_blank"
+                />
-                  rel="noreferrer"
+                Initialize your client
-                >
-                  here
-                </a>
               </div>
               <CodeBlock
                 lang="js"
@@ -200,7 +197,8 @@ function formatCode(
   if (inputs && inputs.length > 0) {
     codeTemplate = codeTemplate.replace(
       "inputs: {}",
-      `inputs: ${JSON.stringify(
+      `inputs: ${
+        JSON.stringify(
           Object.fromEntries(
             inputs.map((x) => {
               return [x?.input_id, ""];
@@ -211,7 +209,8 @@ function formatCode(
         )
           .split("\n")
           .map((line, index) => (index === 0 ? line : `    ${line}`)) // Add two spaces indentation except for the first line
-        .join("\n")}`,
+          .join("\n")
+      }`,
     );
   } else {
     codeTemplate = codeTemplate.replace(

web/src/components/InsertModal.tsx

@@ -26,11 +26,12 @@ import type { UnknownKeysParam, ZodObject, ZodRawShape, z } from "zod";
 export function InsertModal<
   K extends ZodRawShape,
   Y extends UnknownKeysParam,
-  Z extends ZodObject<K, Y>
+  Z extends ZodObject<K, Y>,
 >(props: {
   tooltip?: string;
   disabled?: boolean;
-  title: string;
+  title: React.ReactNode;
+  buttonTitle?: React.ReactNode;
   description: string;
   dialogClassName?: string;
   serverAction: (data: z.infer<Z>) => Promise<unknown>;
@@ -54,7 +55,7 @@ export function InsertModal<
                 setOpen(true);
               }}
             >
-              {props.title}
+              {props.buttonTitle ?? props.title}
             </Button>
           </TooltipTrigger>
           <TooltipContent>
@@ -105,7 +106,7 @@ export function InsertModal<
 export function UpdateModal<
   K extends ZodRawShape,
   Y extends UnknownKeysParam,
-  Z extends ZodObject<K, Y>
+  Z extends ZodObject<K, Y>,
 >(props: {
   open?: boolean;
   setOpen?: (open: boolean) => void;
@@ -118,7 +119,7 @@ export function UpdateModal<
   serverAction: (
     data: z.infer<Z> & {
       id: string;
-    }
+    },
   ) => Promise<unknown>;
   formSchema: Z;
   fieldConfig?: FieldConfig<z.infer<Z>>;
@@ -165,7 +166,7 @@ export function UpdateModal<
               props.serverAction({
                 ...data,
                 id: props.data.id,
-              })
+              }),
             );
             setIsLoading(false);
             setOpen(false);

web/src/components/MachineList.tsx

@@ -41,6 +41,7 @@ import {
   updateMachine,
 } from "@/server/curdMachine";
 import { editWorkflowOnMachine } from "@/server/editWorkflowOnMachine";
+import { getCurrentPlanWithAuth } from "@/server/getCurrentPlan";
 import type {
   ColumnDef,
   ColumnFiltersState,
@@ -55,7 +56,7 @@ import {
   getSortedRowModel,
   useReactTable,
 } from "@tanstack/react-table";
-import { ArrowUpDown, MoreHorizontal } from "lucide-react";
+import { ArrowUpDown, Lock, MoreHorizontal, Plus } from "lucide-react";
 import * as React from "react";
 import { useState } from "react";
 import { toast } from "sonner";
@@ -148,9 +149,16 @@ export const columns: ColumnDef<Machine>[] = [
     header: () => <div className="text-left">Type</div>,
     cell: ({ row }) => {
       return (
-        <div className="text-left font-medium truncate">
+        <Badge
+          className="text-left font-medium truncate"
+          variant={
+            row.original.type == "comfy-deploy-serverless"
+              ? "success"
+              : "outline"
+          }
+        >
           {row.original.type}
-        </div>
+        </Badge>
       );
     },
   },
@@ -182,6 +190,7 @@ export const columns: ColumnDef<Machine>[] = [
     cell: ({ row }) => {
       const machine = row.original;
       const [open, setOpen] = useState(false);
+      const sub = useCurrentPlan();
 
       return (
         <DropdownMenu>
@@ -290,7 +299,10 @@ export const columns: ColumnDef<Machine>[] = [
                   fieldType: "models",
                 },
                 gpu: {
-                  inputProps: {},
+                  fieldType: "gpuPicker",
+                  inputProps: {
+                    sub: sub,
+                  },
                 },
               }}
             />
@@ -318,6 +330,8 @@ export const columns: ColumnDef<Machine>[] = [
   },
 ];
 
+import { useCurrentPlan } from "./useCurrentPlan";
+
 export function MachineList({
   data,
   userMetadata,
@@ -333,6 +347,8 @@ export function MachineList({
     React.useState<VisibilityState>({});
   const [rowSelection, setRowSelection] = React.useState({});
 
+  const sub = useCurrentPlan();
+
   const table = useReactTable({
     data,
     columns,
@@ -352,6 +368,21 @@ export function MachineList({
     },
   });
 
+  let machineMaxCount = 2;
+
+  // Temp fixes for machine count
+  if (userMetadata.betaFeaturesAccess) machineMaxCount = 5;
+
+  if (sub?.plan == "pro") {
+    machineMaxCount = 10;
+  } else if (sub?.plan == "enterprise") {
+    machineMaxCount = 99;
+  }
+
+  const locked =
+    data.some((machine) => machine.type === "modal-serverless") &&
+    data.length >= machineMaxCount;
+
   return (
     <div className="w-full">
       <div className="flex items-center py-4">
@@ -366,17 +397,18 @@ export function MachineList({
         <div className="ml-auto flex gap-2">
           <InsertModal
             dialogClassName="sm:max-w-[600px]"
-            disabled={
+            disabled={locked}
-              data.some(
-                (machine) => machine.type === "comfy-deploy-serverless",
-              ) && !userMetadata.betaFeaturesAccess
-            }
             tooltip={
-              data.some((machine) => machine.type === "comfy-deploy-serverless")
+              locked
-                ? "Only one hosted machine at preview stage"
+                ? `Max ${machineMaxCount} ComfyUI machine for your account, upgrade to unlock more cnfiguration.`
-                : undefined
+                : `Max ${machineMaxCount} ComfyUI machine for your account`
             }
-            title="New Machine"
+            buttonTitle={
+              <>
+                New Machine {locked ? <Lock size={14} /> : <Plus size={14} />}
+              </>
+            }
+            title={"New Machine"}
             description="Add custom ComfyUI machines to your account."
             serverAction={addCustomMachine}
             formSchema={addCustomMachineSchema}
@@ -402,11 +434,9 @@ export function MachineList({
                 },
               },
               gpu: {
-                fieldType: !userMetadata.betaFeaturesAccess
+                fieldType: "gpuPicker",
-                  ? "fallback"
-                  : "select",
                 inputProps: {
-                  disabled: !userMetadata.betaFeaturesAccess,
+                  sub: sub,
                 },
               },
             }}

web/src/components/ModelList.tsx

@@ -15,7 +15,7 @@ import {
   TableHeader,
   TableRow,
 } from "@/components/ui/table";
-import type { getAllUserModels as getAllUserModels } from "@/server/getAllUserModel";
+import type { getAllUserModels } from "@/server/getAllUserModel";
 import type {
   ColumnDef,
   ColumnFiltersState,
@@ -32,8 +32,8 @@ import {
 } from "@tanstack/react-table";
 import { ArrowUpDown } from "lucide-react";
 import * as React from "react";
-import { addCivitaiModel } from "@/server/curdModel";
+import { addModel } from "@/server/curdModel";
-import { addCivitaiModelSchema } from "@/server/addCivitaiModelSchema";
+import { downloadUrlModelSchema } from "@/server/addCivitaiModelSchema";
 import { modelEnumType } from "@/db/schema";
 
 export type ModelItemList = NonNullable<
@@ -46,8 +46,10 @@ export const columns: ColumnDef<ModelItemList>[] = [
     id: "select",
     header: ({ table }) => (
       <Checkbox
-        checked={table.getIsAllPageRowsSelected() ||
+        checked={
-          (table.getIsSomePageRowsSelected() && "indeterminate")}
+          table.getIsAllPageRowsSelected() ||
+          (table.getIsSomePageRowsSelected() && "indeterminate")
+        }
         onCheckedChange={(value) => table.toggleAllPageRowsSelected(!!value)}
         aria-label="Select all"
       />
@@ -79,19 +81,15 @@ export const columns: ColumnDef<ModelItemList>[] = [
       const model = row.original;
       return (
         <>
-          {
+          {/*<a
-            /*<a
           className="hover:underline flex gap-2"
           href={`/storage/${model.id}`} // TODO
-        >*/
+        >*/}
-          }
           <span className="truncate max-w-[200px]">
             {row.original.model_name}
           </span>
 
-          {model.is_public
+          {model.is_public ? <></> : <Badge variant="orange">Private</Badge>}
-            ? <></>
-            : <Badge variant="orange">Private</Badge>}
         </>
       );
     },
@@ -112,9 +110,13 @@ export const columns: ColumnDef<ModelItemList>[] = [
     cell: ({ row }) => {
       return (
         <Badge
-          variant={row.original.status === "failed"
+          variant={
+            row.original.status === "failed"
               ? "red"
-            : (row.original.status === "started" ? "yellow" : "green")}
+              : row.original.status === "started"
+                ? "yellow"
+                : "green"
+          }
         >
           {row.original.status}
         </Badge>
@@ -186,10 +188,10 @@ export const columns: ColumnDef<ModelItemList>[] = [
     },
     cell: ({ row }) => {
       const model_type_map: Record<modelEnumType, any> = {
-        "checkpoint": "amber",
+        checkpoint: "amber",
-        "lora": "green",
+        lora: "green",
-        "embedding": "violet",
+        embedding: "violet",
-        "vae": "teal",
+        vae: "teal",
       };
 
       function getBadgeColor(modelType: modelEnumType) {
@@ -259,9 +261,8 @@ export function ModelList({ data }: { data: ModelItemList[] }) {
   const [columnFilters, setColumnFilters] = React.useState<ColumnFiltersState>(
     [],
   );
-  const [columnVisibility, setColumnVisibility] = React.useState<
+  const [columnVisibility, setColumnVisibility] =
-    VisibilityState
+    React.useState<VisibilityState>({});
-  >({});
   const [rowSelection, setRowSelection] = React.useState({});
 
   const table = useReactTable({
@@ -288,27 +289,27 @@ export function ModelList({ data }: { data: ModelItemList[] }) {
       <div className="flex flex-row w-full items-center py-4">
         <Input
           placeholder="Filter workflows..."
-          value={(table.getColumn("model_name")?.getFilterValue() as string) ??
+          value={
-            ""}
+            (table.getColumn("model_name")?.getFilterValue() as string) ?? ""
+          }
           onChange={(event) =>
-            table.getColumn("model_name")?.setFilterValue(event.target.value)}
+            table.getColumn("model_name")?.setFilterValue(event.target.value)
+          }
           className="max-w-sm"
         />
         <div className="ml-auto flex gap-2">
           <InsertModal
             dialogClassName="sm:max-w-[600px]"
             disabled={
-              false
+              false // TODO: limitations based on plan
-              // TODO: limitations based on plan
             }
-            tooltip={"Add models using their civitai url!"}
+            title="Add a Model"
-            title="Add a Civitai Model"
+            description="using a link to a model"
-            description="Pick a model from civitai"
+            serverAction={addModel}
-            serverAction={addCivitaiModel}
+            formSchema={downloadUrlModelSchema}
-            formSchema={addCivitaiModelSchema}
             fieldConfig={{
-              civitai_url: {
+              url: {
-                fieldType: "fallback",
+                fieldType: "modelUrlPicker",
                 inputProps: { required: true },
                 description: (
                   <>
@@ -317,13 +318,21 @@ export function ModelList({ data }: { data: ModelItemList[] }) {
                       href="https://www.civitai.com/models"
                       target="_blank"
                       className="underline text-blue-600 hover:text-blue-800 visited:text-purple-600"
+                      rel="noreferrer"
                     >
                       civitai.com
                     </a>{" "}
-                    and place it's url here
+                    or a url we can download a model from
                   </>
                 ),
               },
+              model_type: {
+                fieldType: "select",
+                inputProps: { required: true },
+                description: (
+                  <>We'll figure this out if you pick a civitai model</>
+                ),
+              },
             }}
           />
         </div>
@@ -336,7 +345,9 @@ export function ModelList({ data }: { data: ModelItemList[] }) {
                 {headerGroup.headers.map((header) => {
                   return (
                     <TableHead key={header.id}>
-                      {header.isPlaceholder ? null : flexRender(
+                      {header.isPlaceholder
+                        ? null
+                        : flexRender(
                             header.column.columnDef.header,
                             header.getContext(),
                           )}
@@ -347,8 +358,7 @@ export function ModelList({ data }: { data: ModelItemList[] }) {
             ))}
           </TableHeader>
           <TableBody>
-            {table.getRowModel().rows?.length
+            {table.getRowModel().rows?.length ? (
-              ? (
               table.getRowModel().rows.map((row) => (
                 <TableRow
                   key={row.id}
@@ -364,8 +374,7 @@ export function ModelList({ data }: { data: ModelItemList[] }) {
                   ))}
                 </TableRow>
               ))
-              )
+            ) : (
-              : (
               <TableRow>
                 <TableCell
                   colSpan={columns.length}

web/src/components/VersionSelect.tsx

@@ -50,7 +50,7 @@ import {
   Play,
 } from "lucide-react";
 import { parseAsInteger, useQueryState } from "next-usequerystate";
-import { useEffect, useMemo, useState } from "react";
+import { useCallback, useEffect, useMemo, useState } from "react";
 import { toast } from "sonner";
 import useSWR from "swr";
 import type { z } from "zod";
@@ -60,6 +60,7 @@ import { callServerPromise } from "./callServerPromise";
 import fetcher from "./fetcher";
 import { ButtonAction } from "@/components/ButtonActionLoader";
 import { editWorkflowOnMachine } from "@/server/editWorkflowOnMachine";
+import { usePathname, useRouter, useSearchParams } from "next/navigation";
 
 export function VersionSelect({
   workflow,
@@ -126,12 +127,27 @@ export function MachineSelect({
 
 export function useSelectedMachine(
   machines: Awaited<ReturnType<typeof getMachines>>,
-) {
+): [string, (v: string) => void] {
-  const a = useQueryState("machine", {
+  const searchParams = useSearchParams();
-    defaultValue: machines?.[0]?.id ?? "",
+  const pathname = usePathname();
-  });
+  const router = useRouter();
 
-  return a;
+  const createQueryString = useCallback(
+    (name: string, value: string) => {
+      const params = new URLSearchParams(searchParams.toString());
+      params.set(name, value);
+
+      return params.toString();
+    },
+    [searchParams],
+  );
+
+  return [
+    searchParams.get("machine") ?? machines?.[0]?.id ?? "",
+    (v: string) => {
+      router.push(pathname + "?" + createQueryString("machine", v));
+    },
+  ];
 }
 
 type PublicRunStore = {

web/src/components/custom-form/CivitaiModelRegistry.tsx

@@ -0,0 +1,72 @@
+"use client";
+import type { AutoFormInputComponentProps } from "../ui/auto-form/types";
+import * as React from "react";
+import { useDebouncedCallback } from "use-debounce";
+import { z } from "zod";
+import { CivitalModelSchema, ModelListWrapper } from "./CivitalModelSchema";
+import { getUrl, mapModelsList } from "./getUrl";
+import { ModelSelector } from "./ModelSelector";
+
+export function CivitaiModelRegistry({
+  field,
+  selectMultiple = true,
+}: Pick<AutoFormInputComponentProps, "field"> & {
+  selectMultiple?: boolean;
+}) {
+  const [modelList, setModelList] =
+    React.useState<z.infer<typeof ModelListWrapper>>();
+
+  const [loading, setLoading] = React.useState(false);
+
+  const handleSearch = useDebouncedCallback((search) => {
+    console.log(`Searching... ${search}`);
+
+    setLoading(true);
+
+    const controller = new AbortController();
+    fetch(getUrl(search), {
+      signal: controller.signal,
+    })
+      .then((x) => x.json())
+      .then((a) => {
+        const list = CivitalModelSchema.parse(a);
+        console.log(a);
+
+        setModelList(mapModelsList(list));
+        setLoading(false);
+      });
+
+    return () => {
+      controller.abort();
+      setLoading(false);
+    };
+  }, 300);
+
+  React.useEffect(() => {
+    const controller = new AbortController();
+    fetch(getUrl(), {
+      signal: controller.signal,
+    })
+      .then((x) => x.json())
+      .then((a) => {
+        const list = CivitalModelSchema.parse(a);
+        setModelList(mapModelsList(list));
+      });
+
+    return () => {
+      controller.abort();
+    };
+  }, []);
+
+  return (
+    <ModelSelector
+      selectMultiple={selectMultiple}
+      field={field}
+      modelList={modelList}
+      label="Civitai"
+      onSearch={handleSearch}
+      shouldFilter={false}
+      isLoading={loading}
+    />
+  );
+}

web/src/components/custom-form/CivitalModelSchema.tsx

@@ -0,0 +1,94 @@
+"use client";
+import { z } from "zod";
+
+export const Model = z.object({
+  name: z.string(),
+  type: z.string(),
+  base: z.string(),
+  save_path: z.string(),
+  description: z.string(),
+  reference: z.string(),
+  filename: z.string(),
+  url: z.string(),
+});
+
+export const CivitalModelSchema = z.object({
+  items: z.array(
+    z.object({
+      id: z.number(),
+      name: z.string(),
+      description: z.string(),
+      type: z.string(),
+      creator: z
+        .object({
+          username: z.string().nullable(),
+          image: z.string().nullable().default(null),
+        })
+        .nullable(),
+      tags: z.array(z.string()),
+      modelVersions: z.array(
+        z.object({
+          id: z.number(),
+          modelId: z.number(),
+          name: z.string(),
+          createdAt: z.string(),
+          updatedAt: z.string(),
+          status: z.string(),
+          publishedAt: z.string(),
+          trainedWords: z.array(z.unknown()),
+          trainingStatus: z.string().nullable(),
+          trainingDetails: z.string().nullable(),
+          baseModel: z.string(),
+          baseModelType: z.string().nullable(),
+          earlyAccessTimeFrame: z.number(),
+          description: z.string().nullable(),
+          vaeId: z.number().nullable(),
+          stats: z.object({
+            downloadCount: z.number(),
+            ratingCount: z.number(),
+            rating: z.number(),
+          }),
+          files: z.array(
+            z.object({
+              id: z.number(),
+              sizeKB: z.number(),
+              name: z.string(),
+              type: z.string(),
+              downloadUrl: z.string(),
+            }),
+          ),
+          images: z.array(
+            z.object({
+              id: z.number(),
+              url: z.string(),
+              nsfw: z.string(),
+              width: z.number(),
+              height: z.number(),
+              hash: z.string(),
+              type: z.string(),
+              metadata: z.object({
+                hash: z.string(),
+                width: z.number(),
+                height: z.number(),
+              }),
+              meta: z.any(),
+            }),
+          ),
+          downloadUrl: z.string(),
+        }),
+      ),
+    }),
+  ),
+  metadata: z.object({
+    totalItems: z.number(),
+    currentPage: z.number(),
+    pageSize: z.number(),
+    totalPages: z.number(),
+    nextPage: z.string().optional(),
+  }),
+});
+export const ModelList = z.array(Model);
+
+export const ModelListWrapper = z.object({
+  models: ModelList,
+});

web/src/components/custom-form/ComfyUIManagerModelRegistry.tsx

@@ -0,0 +1,43 @@
+"use client";
+import type { AutoFormInputComponentProps } from "../ui/auto-form/types";
+import * as React from "react";
+import { z } from "zod";
+import { ModelListWrapper } from "./CivitalModelSchema";
+import { ModelSelector } from "./ModelSelector";
+
+export function ComfyUIManagerModelRegistry({
+  field,
+  selectMultiple = true,
+}: Pick<AutoFormInputComponentProps, "field"> & {
+  selectMultiple?: boolean;
+}) {
+  const [modelList, setModelList] =
+    React.useState<z.infer<typeof ModelListWrapper>>();
+
+  React.useEffect(() => {
+    const controller = new AbortController();
+    fetch(
+      "https://raw.githubusercontent.com/ltdrdata/ComfyUI-Manager/main/model-list.json",
+      {
+        signal: controller.signal,
+      },
+    )
+      .then((x) => x.json())
+      .then((a) => {
+        setModelList(ModelListWrapper.parse(a));
+      });
+
+    return () => {
+      controller.abort();
+    };
+  }, []);
+
+  return (
+    <ModelSelector
+      selectMultiple={selectMultiple}
+      field={field}
+      modelList={modelList}
+      label="ComfyUI Manager"
+    />
+  );
+}

web/src/components/custom-form/ModelPickerView.tsx

@@ -1,160 +1,17 @@
 "use client";
 
 import type { AutoFormInputComponentProps } from "../ui/auto-form/types";
-import { LoadingIcon } from "@/components/LoadingIcon";
 import {
   Accordion,
   AccordionContent,
   AccordionItem,
   AccordionTrigger,
 } from "@/components/ui/accordion";
-import { Button } from "@/components/ui/button";
-import {
-  Command,
-  CommandEmpty,
-  CommandGroup,
-  CommandInput,
-  CommandItem,
-  CommandList,
-} from "@/components/ui/command";
-import {
-  Popover,
-  PopoverContent,
-  PopoverTrigger,
-} from "@/components/ui/popover";
 import { ScrollArea } from "@/components/ui/scroll-area";
 import { Textarea } from "@/components/ui/textarea";
-import { cn } from "@/lib/utils";
-import { Check, ChevronsUpDown } from "lucide-react";
 import * as React from "react";
-import { useRef } from "react";
+import { CivitaiModelRegistry } from "./CivitaiModelRegistry";
-import { useDebouncedCallback } from "use-debounce";
+import { ComfyUIManagerModelRegistry } from "./ComfyUIManagerModelRegistry";
-import { z } from "zod";
-
-const Model = z.object({
-  name: z.string(),
-  type: z.string(),
-  base: z.string(),
-  save_path: z.string(),
-  description: z.string(),
-  reference: z.string(),
-  filename: z.string(),
-  url: z.string(),
-});
-
-export const CivitalModelSchema = z.object({
-  items: z.array(
-    z.object({
-      id: z.number(),
-      name: z.string(),
-      description: z.string(),
-      type: z.string(),
-      // poi: z.boolean(),
-      // nsfw: z.boolean(),
-      // allowNoCredit: z.boolean(),
-      // allowCommercialUse: z.string(),
-      // allowDerivatives: z.boolean(),
-      // allowDifferentLicense: z.boolean(),
-      // stats: z.object({
-      //   downloadCount: z.number(),
-      //   favoriteCount: z.number(),
-      //   commentCount: z.number(),
-      //   ratingCount: z.number(),
-      //   rating: z.number(),
-      //   tippedAmountCount: z.number(),
-      // }),
-      creator: z
-        .object({
-          username: z.string().nullable(),
-          image: z.string().nullable().default(null),
-        })
-        .nullable(),
-      tags: z.array(z.string()),
-      modelVersions: z.array(
-        z.object({
-          id: z.number(),
-          modelId: z.number(),
-          name: z.string(),
-          createdAt: z.string(),
-          updatedAt: z.string(),
-          status: z.string(),
-          publishedAt: z.string(),
-          trainedWords: z.array(z.unknown()),
-          trainingStatus: z.string().nullable(),
-          trainingDetails: z.string().nullable(),
-          baseModel: z.string(),
-          baseModelType: z.string().nullable(),
-          earlyAccessTimeFrame: z.number(),
-          description: z.string().nullable(),
-          vaeId: z.number().nullable(),
-          stats: z.object({
-            downloadCount: z.number(),
-            ratingCount: z.number(),
-            rating: z.number(),
-          }),
-          files: z.array(
-            z.object({
-              id: z.number(),
-              sizeKB: z.number(),
-              name: z.string(),
-              type: z.string(),
-              // metadata: z.object({
-              //   fp: z.string().nullable().optional(),
-              //   size: z.string().nullable().optional(),
-              //   format: z.string().nullable().optional(),
-              // }),
-              // pickleScanResult: z.string(),
-              // pickleScanMessage: z.string(),
-              // virusScanResult: z.string(),
-              // virusScanMessage: z.string().nullable(),
-              // scannedAt: z.string(),
-              // hashes: z.object({
-              //   AutoV1: z.string().nullable().optional(),
-              //   AutoV2: z.string().nullable().optional(),
-              //   SHA256: z.string().nullable().optional(),
-              //   CRC32: z.string().nullable().optional(),
-              //   BLAKE3: z.string().nullable().optional(),
-              // }),
-              downloadUrl: z.string(),
-              // primary: z.boolean().default(false),
-            })
-          ),
-          images: z.array(
-            z.object({
-              id: z.number(),
-              url: z.string(),
-              nsfw: z.string(),
-              width: z.number(),
-              height: z.number(),
-              hash: z.string(),
-              type: z.string(),
-              metadata: z.object({
-                hash: z.string(),
-                width: z.number(),
-                height: z.number(),
-              }),
-              meta: z.any(),
-            })
-          ),
-          downloadUrl: z.string(),
-        })
-      ),
-    })
-  ),
-  metadata: z.object({
-    totalItems: z.number(),
-    currentPage: z.number(),
-    pageSize: z.number(),
-    totalPages: z.number(),
-    nextPage: z.string().optional(),
-  }),
-});
-
-const ModelList = z.array(Model);
-
-export const ModelListWrapper = z.object({
-  models: ModelList,
-});
 
 export function ModelPickerView({
   field,
@@ -187,240 +44,3 @@ export function ModelPickerView({
     </Accordion>
   );
 }
-
-function mapType(type: string) {
-  switch (type) {
-    case "checkpoint":
-      return "checkpoints";
-  }
-  return type;
-}
-
-function mapModelsList(
-  models: z.infer<typeof CivitalModelSchema>
-): z.infer<typeof ModelListWrapper> {
-  return {
-    models: models.items.flatMap((item) => {
-      return item.modelVersions.map((v) => {
-        return {
-          name: `${item.name} ${v.name} (${v.files[0].name})`,
-          type: mapType(item.type.toLowerCase()),
-          base: v.baseModel,
-          save_path: "default",
-          description: item.description,
-          reference: "",
-          filename: v.files[0].name,
-          url: v.files[0].downloadUrl,
-        } as z.infer<typeof Model>;
-      });
-    }),
-  };
-}
-
-function getUrl(search?: string) {
-  const baseUrl = "https://civitai.com/api/v1/models";
-  const searchParams = {
-    limit: 5,
-  } as any;
-  searchParams["sort"] = "Most Downloaded";
-
-  if (search) {
-    searchParams["query"] = search;
-  } else {
-    // sort: "Highest Rated",
-  }
-
-  const url = new URL(baseUrl);
-  Object.keys(searchParams).forEach((key) =>
-    url.searchParams.append(key, searchParams[key])
-  );
-
-  return url;
-}
-
-export function CivitaiModelRegistry({
-  field,
-}: Pick<AutoFormInputComponentProps, "field">) {
-  const [modelList, setModelList] =
-    React.useState<z.infer<typeof ModelListWrapper>>();
-
-  const [loading, setLoading] = React.useState(false);
-
-  const handleSearch = useDebouncedCallback((search) => {
-    console.log(`Searching... ${search}`);
-
-    setLoading(true);
-
-    const controller = new AbortController();
-    fetch(getUrl(search), {
-      signal: controller.signal,
-    })
-      .then((x) => x.json())
-      .then((a) => {
-        const list = CivitalModelSchema.parse(a);
-        console.log(a);
-
-        setModelList(mapModelsList(list));
-        setLoading(false);
-      });
-
-    return () => {
-      controller.abort();
-      setLoading(false);
-    };
-  }, 300);
-
-  React.useEffect(() => {
-    const controller = new AbortController();
-    fetch(getUrl(), {
-      signal: controller.signal,
-    })
-      .then((x) => x.json())
-      .then((a) => {
-        const list = CivitalModelSchema.parse(a);
-        setModelList(mapModelsList(list));
-      });
-
-    return () => {
-      controller.abort();
-    };
-  }, []);
-
-  return (
-    <ModelSelector
-      field={field}
-      modelList={modelList}
-      label="Civitai"
-      onSearch={handleSearch}
-      shouldFilter={false}
-      isLoading={loading}
-    />
-  );
-}
-
-export function ComfyUIManagerModelRegistry({
-  field,
-}: Pick<AutoFormInputComponentProps, "field">) {
-  const [modelList, setModelList] =
-    React.useState<z.infer<typeof ModelListWrapper>>();
-
-  React.useEffect(() => {
-    const controller = new AbortController();
-    fetch(
-      "https://raw.githubusercontent.com/ltdrdata/ComfyUI-Manager/main/model-list.json",
-      {
-        signal: controller.signal,
-      }
-    )
-      .then((x) => x.json())
-      .then((a) => {
-        setModelList(ModelListWrapper.parse(a));
-      });
-
-    return () => {
-      controller.abort();
-    };
-  }, []);
-
-  return (
-    <ModelSelector
-      field={field}
-      modelList={modelList}
-      label="ComfyUI Manager"
-    />
-  );
-}
-
-export function ModelSelector({
-  field,
-  modelList,
-  label,
-  onSearch,
-  shouldFilter = true,
-  isLoading,
-}: Pick<AutoFormInputComponentProps, "field"> & {
-  modelList?: z.infer<typeof ModelListWrapper>;
-  label: string;
-  onSearch?: (search: string) => void;
-  shouldFilter?: boolean;
-  isLoading?: boolean;
-}) {
-  const value = (field.value as z.infer<typeof ModelList>) ?? [];
-  const [open, setOpen] = React.useState(false);
-
-  function toggleModel(model: z.infer<typeof Model>) {
-    const prevSelectedModels = value;
-    if (
-      prevSelectedModels.some(
-        (selectedModel) =>
-          selectedModel.url + selectedModel.name === model.url + model.name
-      )
-    ) {
-      field.onChange(
-        prevSelectedModels.filter(
-          (selectedModel) =>
-            selectedModel.url + selectedModel.name !== model.url + model.name
-        )
-      );
-    } else {
-      field.onChange([...prevSelectedModels, model]);
-    }
-  }
-
-  const containerRef = useRef<HTMLDivElement>(null);
-
-  return (
-    <div className="" ref={containerRef}>
-      <Popover open={open} onOpenChange={setOpen}>
-        <PopoverTrigger asChild>
-          <Button
-            variant="outline"
-            role="combobox"
-            aria-expanded={open}
-            className="w-full justify-between flex"
-          >
-            Add from {label}
-            <ChevronsUpDown className="ml-2 h-4 w-4 shrink-0 opacity-50" />
-          </Button>
-        </PopoverTrigger>
-        <PopoverContent className="w-[375px] p-0" side="bottom">
-          <Command shouldFilter={shouldFilter}>
-            <CommandInput
-              placeholder="Search models..."
-              className="h-9"
-              onValueChange={onSearch}
-            >
-              {isLoading && <LoadingIcon />}
-            </CommandInput>
-            <CommandEmpty>No models found.</CommandEmpty>
-            <CommandList className="pointer-events-auto">
-              <CommandGroup>
-                {modelList?.models.map((model) => (
-                  <CommandItem
-                    key={model.url + model.name}
-                    value={model.url}
-                    onSelect={() => {
-                      toggleModel(model);
-                    }}
-                  >
-                    {model.name}
-                    <Check
-                      className={cn(
-                        "ml-auto h-4 w-4",
-                        value.some(
-                          (selectedModel) => selectedModel.url === model.url
-                        )
-                          ? "opacity-100"
-                          : "opacity-0"
-                      )}
-                    />
-                  </CommandItem>
-                ))}
-              </CommandGroup>
-            </CommandList>
-          </Command>
-        </PopoverContent>
-      </Popover>
-    </div>
-  );
-}

web/src/components/custom-form/ModelSelector.tsx

@@ -0,0 +1,123 @@
+"use client";
+import type { AutoFormInputComponentProps } from "../ui/auto-form/types";
+import { LoadingIcon } from "@/components/LoadingIcon";
+import { Button } from "@/components/ui/button";
+import {
+  Command,
+  CommandEmpty,
+  CommandGroup,
+  CommandInput,
+  CommandItem,
+  CommandList,
+} from "@/components/ui/command";
+import {
+  Popover,
+  PopoverContent,
+  PopoverTrigger,
+} from "@/components/ui/popover";
+import { cn } from "@/lib/utils";
+import { Check, ChevronsUpDown } from "lucide-react";
+import * as React from "react";
+import { useRef } from "react";
+import { z } from "zod";
+import { ModelListWrapper, Model, ModelList } from "./CivitalModelSchema";
+
+export function ModelSelector({
+  field,
+  modelList,
+  label,
+  onSearch,
+  shouldFilter = true,
+  isLoading,
+  selectMultiple = true,
+}: Pick<AutoFormInputComponentProps, "field"> & {
+  modelList?: z.infer<typeof ModelListWrapper>;
+  label: string;
+  onSearch?: (search: string) => void;
+  shouldFilter?: boolean;
+  isLoading?: boolean;
+  selectMultiple?: boolean;
+}) {
+  const value = (field.value as z.infer<typeof ModelList>) ?? [];
+  const [open, setOpen] = React.useState(false);
+
+  function toggleModel(model: z.infer<typeof Model>) {
+    const prevSelectedModels = value;
+    if (
+      prevSelectedModels.some(
+        (selectedModel) =>
+          selectedModel.url + selectedModel.name === model.url + model.name,
+      )
+    ) {
+      field.onChange(
+        prevSelectedModels.filter(
+          (selectedModel) =>
+            selectedModel.url + selectedModel.name !== model.url + model.name,
+        ),
+      );
+    } else {
+      if (!selectMultiple) {
+        field.onChange([model]);
+      } else {
+        field.onChange([...prevSelectedModels, model]);
+      }
+    }
+  }
+
+  const containerRef = useRef<HTMLDivElement>(null);
+
+  return (
+    <div className="" ref={containerRef}>
+      <Popover open={open} onOpenChange={setOpen}>
+        <PopoverTrigger asChild>
+          <Button
+            variant="outline"
+            role="combobox"
+            aria-expanded={open}
+            className="w-full justify-between flex"
+          >
+            Add from {label}
+            <ChevronsUpDown className="ml-2 h-4 w-4 shrink-0 opacity-50" />
+          </Button>
+        </PopoverTrigger>
+        <PopoverContent className="w-[375px] p-0" side="bottom">
+          <Command shouldFilter={shouldFilter}>
+            <CommandInput
+              placeholder="Search models..."
+              className="h-9"
+              onValueChange={onSearch}
+            >
+              {isLoading && <LoadingIcon />}
+            </CommandInput>
+            <CommandEmpty>No models found.</CommandEmpty>
+            <CommandList className="pointer-events-auto">
+              <CommandGroup>
+                {modelList?.models.map((model) => (
+                  <CommandItem
+                    key={model.url + model.name}
+                    value={model.url}
+                    onSelect={() => {
+                      toggleModel(model);
+                    }}
+                  >
+                    {model.name}
+                    <Check
+                      className={cn(
+                        "ml-auto h-4 w-4",
+                        value.some(
+                          (selectedModel) => selectedModel.url === model.url,
+                        )
+                          ? "opacity-100"
+                          : "opacity-0",
+                      )}
+                    />
+                  </CommandItem>
+                ))}
+              </CommandGroup>
+            </CommandList>
+          </Command>
+        </PopoverContent>
+      </Popover>
+    </div>
+  );
+}

web/src/components/custom-form/SnapshotPickerView.tsx

@@ -8,7 +8,9 @@ import {
   AccordionItem,
   AccordionTrigger,
 } from "@/components/ui/accordion";
+import { Badge } from "@/components/ui/badge";
 import { Button } from "@/components/ui/button";
+import { Card } from "@/components/ui/card";
 import {
   Command,
   CommandEmpty,
@@ -17,6 +19,22 @@ import {
   CommandItem,
   CommandList,
 } from "@/components/ui/command";
+import {
+  Dialog,
+  DialogClose,
+  DialogContent,
+  DialogDescription,
+  DialogFooter,
+  DialogHeader,
+  DialogTitle,
+  DialogTrigger,
+} from "@/components/ui/dialog";
+import {
+  DropdownMenu,
+  DropdownMenuContent,
+  DropdownMenuItem,
+  DropdownMenuTrigger,
+} from "@/components/ui/dropdown-menu";
 import {
   Popover,
   PopoverContent,
@@ -25,7 +43,15 @@ import {
 import { Textarea } from "@/components/ui/textarea";
 import { cn } from "@/lib/utils";
 import { findAllDeployments } from "@/server/curdDeploments";
-import { Check, ChevronsUpDown } from "lucide-react";
+import {
+  Check,
+  ChevronsUpDown,
+  Edit,
+  ExternalLink,
+  FolderInput,
+  MoreVertical,
+  Plus,
+} from "lucide-react";
 import * as React from "react";
 import { toast } from "sonner";
 import useSWR from "swr";
@@ -39,19 +65,95 @@ export function SnapshotPickerView({
       <AccordionItem value="item-1">
         <AccordionTrigger className="text-sm">Custom Nodes</AccordionTrigger>
         <AccordionContent className="flex gap-2 flex-col px-1">
-          <SnapshotPresetPicker field={field} />
+          <div className="flex flex-wrap gap-2 justify-end">
             <CustomNodesSelector field={field} />
-          {field.value && (
+            <SnapshotPresetPicker field={field} />
-            // <ScrollArea className="w-full bg-gray-100 mx-auto max-w-[500px] rounded-lg">
+            <Dialog>
+              <DialogTrigger asChild>
+                <Button variant="secondary" className="w-fit">
+                  Edit <Edit size={14}></Edit>
+                </Button>
+              </DialogTrigger>
+              <DialogContent className="sm:max-w-[600px] h-full max-h-[600px] grid grid-rows-[auto,1fr,auto]">
+                <DialogHeader>
+                  <DialogTitle>Edit custom nodes</DialogTitle>
+                  <DialogDescription>
+                    Make advacne changes to the snapshots
+                  </DialogDescription>
+                </DialogHeader>
                 <Textarea
-              className="min-h-[150px] max-h-[300px] p-2 rounded-md text-xs w-full"
+                  className="h-full p-2 max-h-[600px] rounded-md text-xs w-full"
                   value={JSON.stringify(field.value, null, 2)}
                   onChange={(e) => {
                     // Update field.onChange to pass the array of selected models
                     field.onChange(JSON.parse(e.target.value));
                   }}
                 />
-            // </ScrollArea>
+                <DialogFooter>
+                  <DialogClose>
+                    <Button type="button" variant="secondary">
+                      Close
+                    </Button>
+                  </DialogClose>
+                </DialogFooter>
+              </DialogContent>
+            </Dialog>
+          </div>
+          {field.value && (
+            <div className="flex gap-2 flex-col">
+              {Object.entries(field.value.git_custom_nodes).map(
+                ([key, item]: [string, any], index) => (
+                  <Card className="p-4 flex gap-4 items-center justify-between">
+                    <div className="flex gap-4 items-center">
+                      <div className="bg-gray-200 aspect-square w-6 h-6 rounded-full text-center">
+                        {index + 1}
+                      </div>
+                      <div>
+                        <a
+                          target="_blank"
+                          href={key}
+                          className="hover:underline flex items-center gap-2"
+                          rel="noreferrer"
+                        >
+                          <ExternalLink size={12} /> {key}
+                        </a>
+                        <div className="text-2xs text-primary/50">
+                          {item.hash}
+                        </div>
+                      </div>
+                    </div>
+
+                    <DropdownMenu>
+                      <DropdownMenuTrigger asChild type="button">
+                        <Button type="button" variant={"ghost"}>
+                          <MoreVertical size={12} />
+                        </Button>
+                      </DropdownMenuTrigger>
+                      <DropdownMenuContent>
+                        <DropdownMenuItem
+                          disabled={key.endsWith("comfyui-deploy.git")}
+                          // className="opacity-50"
+                          onClick={() => {
+                            const newNodeList = {
+                              ...field.value.git_custom_nodes,
+                            };
+                            delete newNodeList[key];
+                            const nodeList = newNodeList;
+                            const newValue = {
+                              ...field.value,
+                              git_custom_nodes: nodeList,
+                            };
+                            field.onChange(newValue);
+                          }}
+                        >
+                          Delete
+                        </DropdownMenuItem>
+                      </DropdownMenuContent>
+                    </DropdownMenu>
+                  </Card>
+                ),
+              )}
+            </div>
           )}
         </AccordionContent>
       </AccordionItem>
@@ -65,7 +167,8 @@ function SnapshotPresetPicker({
   const [open, setOpen] = React.useState(false);
   const [selected, setSelected] = React.useState<string | null>(null);
 
-  const [frameworks, setFramework] = React.useState<
+  const [frameworks, setFramework] =
+    React.useState<
       {
         id: string;
         label: string;
@@ -108,12 +211,13 @@ function SnapshotPresetPicker({
           variant="outline"
           role="combobox"
           aria-expanded={open}
-          className="w-full justify-between flex"
+          className="w-fit justify-between flex"
         >
-          {selected
+          <FolderInput size={14} />
+          Import
+          {/* {selected
             ? findItem(selected)?.label
-            : "Select snapshot (From deployments)"}
+            : "Select snapshot (From deployments)"} */}
-          <ChevronsUpDown className="ml-2 h-4 w-4 shrink-0 opacity-50" />
         </Button>
       </PopoverTrigger>
       <PopoverContent className="w-[375px] p-0">
@@ -140,7 +244,7 @@ function SnapshotPresetPicker({
                     "ml-auto h-4 w-4",
                     field.value === framework.value
                       ? "opacity-100"
-                      : "opacity-0"
+                      : "opacity-0",
                   )}
                 />
               </CommandItem>
@@ -199,12 +303,12 @@ function CustomNodesSelector({
 
   const { data, error, isLoading } = useSWR<CustomNodeList>(
     "https://raw.githubusercontent.com/ltdrdata/ComfyUI-Manager/main/custom-node-list.json",
-    fetcher
+    fetcher,
   );
 
   const keys = React.useMemo(
     () => Object.keys(customNodeList),
-    [customNodeList, data]
+    [customNodeList, data],
   );
 
   function findItem(value: string) {
@@ -213,6 +317,11 @@ function CustomNodesSelector({
     return included;
   }
 
+  const onChangeRef = React.useRef(field.onChange);
+  React.useEffect(() => {
+    onChangeRef.current = field.onChange;
+  }, [field.onChange]);
+
   return (
     <Popover open={open} onOpenChange={setOpen}>
       <PopoverTrigger asChild>
@@ -220,10 +329,10 @@ function CustomNodesSelector({
           variant="outline"
           role="combobox"
           aria-expanded={open}
-          className="w-full justify-between flex"
+          className="w-fit justify-between flex"
         >
-          Add custom nodes - {keys.length} selected
+          <Plus size={14}></Plus> <Badge>{keys.length} </Badge>
-          <ChevronsUpDown className="ml-2 h-4 w-4 shrink-0 opacity-50" />
+          {/* <ChevronsUpDown className="ml-2 h-4 w-4 shrink-0 opacity-50" /> */}
         </Button>
       </PopoverTrigger>
       <PopoverContent className="w-[375px] p-0" side="bottom">
@@ -243,7 +352,7 @@ function CustomNodesSelector({
                         {
                           hash: string;
                           disabled: boolean;
-                          pip?: string[]
+                          pip?: string[];
                         }
                       >;
                       const x = customNodeList;
@@ -256,7 +365,7 @@ function CustomNodesSelector({
                         const repoName = extractRepoName(currentValue);
                         const id = toast.loading(`Fetching repo info...`);
                         const repo = await fetch(
-                          `https://api.github.com/repos/${repoName}`
+                          `https://api.github.com/repos/${repoName}`,
                         )
                           .then((x) => x.json())
                           .then((x) => {
@@ -267,14 +376,16 @@ function CustomNodesSelector({
                           .catch((e) => {
                             console.error(e);
                             toast.dismiss(id);
-                            toast.error(`Failed to fetch repo info ${e.message}`);
+                            toast.error(
+                              `Failed to fetch repo info ${e.message}`,
+                            );
                             return null;
                           });
 
                         if (!repo) return;
                         const branch = repo.default_branch;
                         const branchInfo = await fetch(
-                          `https://api.github.com/repos/${repoName}/branches/${branch}`
+                          `https://api.github.com/repos/${repoName}/branches/${branch}`,
                         )
                           .then((x) => x.json())
                           .then((x) => BranchInfoSchema.parse(x))
@@ -282,7 +393,7 @@ function CustomNodesSelector({
                             console.error(e);
                             toast.dismiss(id);
                             toast.error(
-                              `Failed to fetch branch info ${e.message}`
+                              `Failed to fetch branch info ${e.message}`,
                             );
                             return null;
                           });
@@ -291,19 +402,27 @@ function CustomNodesSelector({
 
                         if (!branchInfo) return;
 
-                        nodeList = {
+                        const value: Record<string, any> = {
-                          [currentValue]: {
                           hash: branchInfo?.commit.sha,
                           disabled: false,
-                            pip: framework.pip
+                        };
-                          },
+
+                        if (framework.pip) {
+                          value["pip"] = framework.pip;
+                        }
+
+                        nodeList = {
                           ...x,
+                          [currentValue]: value,
                         };
                       }
-                      field.onChange({
+
+                      const newValue = {
                         ...field.value,
                         git_custom_nodes: nodeList,
-                      });
+                      };
+
+                      field.onChange(newValue);
                     }}
                   >
                     {framework.title}
@@ -312,7 +431,7 @@ function CustomNodesSelector({
                         "ml-auto h-4 w-4",
                         findItem(framework.reference)
                           ? "opacity-100"
-                          : "opacity-0"
+                          : "opacity-0",
                       )}
                     />
                   </CommandItem>

web/src/components/custom-form/getUrl.tsx

@@ -0,0 +1,56 @@
+"use client";
+import { z } from "zod";
+import {
+  CivitalModelSchema,
+  ModelListWrapper,
+  Model,
+} from "./CivitalModelSchema";
+
+function mapType(type: string) {
+  switch (type) {
+    case "checkpoint":
+      return "checkpoints";
+  }
+  return type;
+}
+export function mapModelsList(
+  models: z.infer<typeof CivitalModelSchema>,
+): z.infer<typeof ModelListWrapper> {
+  return {
+    models: models.items.flatMap((item) => {
+      return item.modelVersions.map((v) => {
+        return {
+          name: `${item.name} ${v.name} (${v.files[0].name})`,
+          type: mapType(item.type.toLowerCase()),
+          base: v.baseModel,
+          save_path: "default",
+          description: item.description,
+          reference: "",
+          filename: v.files[0].name,
+          // Quick hack to get the download url back as normal url
+          url: `https://civitai.com/models/${v.modelId}?modelVersionId=${v.id}`, //v.files[0].downloadUrl,
+        } as z.infer<typeof Model>;
+      });
+    }),
+  };
+}
+export function getUrl(search?: string) {
+  const baseUrl = "https://civitai.com/api/v1/models";
+  const searchParams = {
+    limit: 5,
+  } as any;
+  searchParams["sort"] = "Most Downloaded";
+
+  if (search) {
+    searchParams["query"] = search;
+  } else {
+    // sort: "Highest Rated",
+  }
+
+  const url = new URL(baseUrl);
+  Object.keys(searchParams).forEach((key) =>
+    url.searchParams.append(key, searchParams[key]),
+  );
+
+  return url;
+}

web/src/components/custom-form/gpu-picker.tsx

@@ -0,0 +1,98 @@
+import { AutoFormInputComponentProps } from "@/components/ui/auto-form/types";
+import { getBaseSchema } from "@/components/ui/auto-form/utils";
+import { Badge } from "@/components/ui/badge";
+import {
+  FormItem,
+  FormLabel,
+  FormControl,
+  FormDescription,
+  FormMessage,
+} from "@/components/ui/form";
+import {
+  Select,
+  SelectContent,
+  SelectItem,
+  SelectTrigger,
+  SelectValue,
+} from "@/components/ui/select";
+import { Lock } from "lucide-react";
+import * as z from "zod";
+
+export default function AutoFormGPUPicker({
+  label,
+  isRequired,
+  field,
+  fieldConfigItem,
+  zodItem,
+}: AutoFormInputComponentProps) {
+  const baseValues = (getBaseSchema(zodItem) as unknown as z.ZodEnum<any>)._def
+    .values;
+
+  let values: [string, string][] = [];
+  if (!Array.isArray(baseValues)) {
+    values = Object.entries(baseValues);
+  } else {
+    values = baseValues.map((value) => [value, value]);
+  }
+
+  function findItem(value: any) {
+    return values.find((item) => item[0] === value);
+  }
+
+  const plan = fieldConfigItem.inputProps?.sub?.plan;
+  const enabledGPU = ["T4"];
+
+  const planMapping: Record<string, string> = {
+    A10G: "pro",
+    A100: "enterprise",
+  };
+
+  if (plan == "pro") {
+    enabledGPU.push("A10G");
+  } else if (plan == "enterprise") {
+    enabledGPU.push("A10G");
+    enabledGPU.push("A100");
+  }
+
+  return (
+    <FormItem>
+      <FormLabel>
+        {label}
+        {isRequired && <span className="text-destructive"> *</span>}
+      </FormLabel>
+      <FormControl>
+        <Select onValueChange={field.onChange} defaultValue={field.value}>
+          <SelectTrigger>
+            <SelectValue
+              className="w-full"
+              placeholder={fieldConfigItem.inputProps?.placeholder}
+            >
+              {field.value ? findItem(field.value)?.[1] : "Select an option"}
+            </SelectValue>
+          </SelectTrigger>
+          <SelectContent>
+            {values.map(([value, label]) => {
+              const enabled = enabledGPU.includes(value);
+              return (
+                <SelectItem value={label} key={value} disabled={!enabled}>
+                  {label}
+                  {!enabled && (
+                    <span className="mx-2 inline-flex items-center justify-center gap-2">
+                      <Badge className="capitalize">{planMapping[value]}</Badge>{" "}
+                      plan required
+                      <Lock size={14}></Lock>
+                    </span>
+                  )}
+                </SelectItem>
+              );
+            })}
+          </SelectContent>
+        </Select>
+      </FormControl>
+      {fieldConfigItem.description && (
+        <FormDescription>{fieldConfigItem.description}</FormDescription>
+      )}
+      <FormMessage />
+    </FormItem>
+  );
+}

web/src/components/custom-form/model-input.tsx

@@ -1,86 +0,0 @@
-// NOTE: this is WIP for doing client side validation for civitai model downloading
-import type { AutoFormInputComponentProps } from "../ui/auto-form/types";
-import { FormControl, FormItem, FormLabel } from "../ui/form";
-import { LoadingIcon } from "@/components/LoadingIcon";
-import * as React from "react";
-import AutoFormInput from "../ui/auto-form/fields/input";
-import { useDebouncedCallback } from "use-debounce";
-import { CivitaiModelResponse } from "@/types/civitai";
-import { z } from "zod";
-import { insertCivitaiModelSchema } from "@/db/schema";
-
-function getUrl(civitai_url: string) {
-  // expect to be a URL to be https://civitai.com/models/36520
-  // possiblity with slugged name and query-param modelVersionId
-
-  const baseUrl = "https://civitai.com/api/v1/models/";
-  const url = new URL(civitai_url);
-  const pathSegments = url.pathname.split("/");
-  const modelId = pathSegments[pathSegments.indexOf("models") + 1];
-  const modelVersionId = url.searchParams.get("modelVersionId");
-
-  return { url: baseUrl + modelId, modelVersionId };
-}
-
-export default function AutoFormCheckpointInput(
-  props: AutoFormInputComponentProps
-) {
-  const [loading, setLoading] = React.useState(false);
-  const [modelRes, setModelRes] =
-    React.useState<z.infer<typeof CivitaiModelResponse>>();
-  const [modelVersionid, setModelVersionId] = React.useState<string | null>();
-  const { label, isRequired, fieldProps, zodItem, fieldConfigItem } = props;
-
-  const handleSearch = useDebouncedCallback((search) => {
-    const validationResult =
-      insertCivitaiModelSchema.shape.civitai_url.safeParse(search);
-    if (!validationResult.success) {
-      console.error(validationResult.error);
-      // Optionally set an error state here
-      return;
-    }
-
-    setLoading(true);
-
-    const controller = new AbortController();
-    const { url, modelVersionId: versionId } = getUrl(search);
-    setModelVersionId(versionId);
-    fetch(url, {
-      signal: controller.signal,
-    })
-      .then((x) => x.json())
-      .then((a) => {
-        const res = CivitaiModelResponse.parse(a);
-        console.log(a);
-        console.log(res);
-        setModelRes(res);
-        setLoading(false);
-      });
-
-    return () => {
-      controller.abort();
-      setLoading(false);
-    };
-  }, 300);
-
-  const modifiedField = {
-    ...fieldProps,
-    // onChange: (event: React.ChangeEvent<HTMLInputElement>) => {
-    //   handleSearch(event.target.value);
-    // },
-  };
-
-  return (
-    <FormItem>
-      {fieldConfigItem.inputProps?.showLabel && (
-        <FormLabel>
-          {label}
-          {isRequired && <span className="text-destructive">*</span>}
-        </FormLabel>
-      )}
-      <FormControl>
-        <AutoFormInput {...props} fieldProps={modifiedField} />
-      </FormControl>
-    </FormItem>
-  );
-}

web/src/components/custom-form/model-picker-url-only.tsx

@@ -0,0 +1,97 @@
+"use client";
+
+import type { AutoFormInputComponentProps } from "../ui/auto-form/types";
+import {
+  FormControl,
+  FormDescription,
+  FormItem,
+  FormLabel,
+  FormMessage,
+} from "../ui/form";
+import { LoadingIcon } from "@/components/LoadingIcon";
+// import { CaretSortIcon, CheckIcon } from "@radix-ui/react-icons";
+import * as React from "react";
+import { Suspense } from "react";
+import {
+  Accordion,
+  AccordionContent,
+  AccordionItem,
+  AccordionTrigger,
+} from "@/components/ui/accordion";
+import { ScrollArea } from "@/components/ui/scroll-area";
+import { Textarea } from "@/components/ui/textarea";
+import { CivitaiModelRegistry } from "./CivitaiModelRegistry";
+import { ComfyUIManagerModelRegistry } from "./ComfyUIManagerModelRegistry";
+import { Input } from "@/components/ui/input";
+import { ModelList } from "@/components/custom-form/CivitalModelSchema";
+import { z } from "zod";
+
+export default function AutoFormModelsPickerUrl({
+  label,
+  isRequired,
+  field,
+  fieldConfigItem,
+  zodItem,
+  fieldProps,
+}: AutoFormInputComponentProps) {
+  return (
+    <FormItem>
+      {fieldConfigItem.inputProps?.showLabel && (
+        <FormLabel>
+          {label}
+          {isRequired && <span className="text-destructive"> *</span>}
+        </FormLabel>
+      )}
+      <FormControl>
+        <Suspense fallback={<LoadingIcon />}>
+          <ModelPickerView field={field} fieldProps={fieldProps} />
+        </Suspense>
+      </FormControl>
+      {fieldConfigItem.description && (
+        <FormDescription>{fieldConfigItem.description}</FormDescription>
+      )}
+      <FormMessage />
+    </FormItem>
+  );
+}
+
+function ModelPickerView({
+  field,
+  fieldProps,
+}: Pick<AutoFormInputComponentProps, "field" | "fieldProps">) {
+  const customOverride = React.useMemo(() => {
+    const customOnChange = (value: z.infer<typeof ModelList>) => {
+      const model = value[0];
+      field.onChange(model?.url);
+    };
+    return {
+      ...field,
+      onChange: customOnChange,
+      value: field.value
+        ? [
+            {
+              url: field.value,
+            },
+          ]
+        : [],
+    };
+  }, [field]);
+
+  return (
+    <div className="flex gap-2 flex-col px-1">
+      <ComfyUIManagerModelRegistry
+        field={customOverride}
+        selectMultiple={false}
+      />
+      <CivitaiModelRegistry field={customOverride} selectMultiple={false} />
+      <Input
+        // className="min-h-[150px] max-h-[300px] p-2 rounded-lg text-xs w-full"
+        value={field.value ?? ""}
+        onChange={(e) => {
+          field.onChange(e.target.value);
+        }}
+        type="text"
+      />
+    </div>
+  );
+}

web/src/components/custom-form/model-picker.tsx

@@ -1,3 +1,5 @@
+"use client";
+
 import type { AutoFormInputComponentProps } from "../ui/auto-form/types";
 import {
   FormControl,
@@ -7,10 +9,19 @@ import {
   FormMessage,
 } from "../ui/form";
 import { LoadingIcon } from "@/components/LoadingIcon";
-import { ModelPickerView } from "@/components/custom-form/ModelPickerView";
 // import { CaretSortIcon, CheckIcon } from "@radix-ui/react-icons";
 import * as React from "react";
 import { Suspense } from "react";
+import {
+  Accordion,
+  AccordionContent,
+  AccordionItem,
+  AccordionTrigger,
+} from "@/components/ui/accordion";
+import { ScrollArea } from "@/components/ui/scroll-area";
+import { Textarea } from "@/components/ui/textarea";
+import { ComfyUIManagerModelRegistry } from "./ComfyUIManagerModelRegistry";
+import { ExternalLink } from "lucide-react";
 
 export default function AutoFormModelsPicker({
   label,
@@ -27,6 +38,7 @@ export default function AutoFormModelsPicker({
           {isRequired && <span className="text-destructive"> *</span>}
         </FormLabel>
       )}
+
       <FormControl>
         <Suspense fallback={<LoadingIcon />}>
           <ModelPickerView field={field} />
@@ -35,7 +47,54 @@ export default function AutoFormModelsPicker({
       {fieldConfigItem.description && (
         <FormDescription>{fieldConfigItem.description}</FormDescription>
       )}
+      <FormDescription>
+        {" "}
+        <div className="text-sm">
+          Models are moving to{" "}
+          <a
+            href="/storage"
+            target="_blank"
+            rel="noreferrer"
+            className="inline-flex items-center gap-1 underline"
+          >
+            <ExternalLink size={12} />
+            Storage
+          </a>
+        </div>
+      </FormDescription>
       <FormMessage />
     </FormItem>
   );
 }
+
+function ModelPickerView({
+  field,
+}: Pick<AutoFormInputComponentProps, "field">) {
+  return (
+    <Accordion type="single" collapsible>
+      <AccordionItem value="item-1">
+        <AccordionTrigger className="text-sm">
+          Models (ComfyUI Manager)
+        </AccordionTrigger>
+        <AccordionContent>
+          <div className="flex gap-2 flex-col px-1">
+            <ComfyUIManagerModelRegistry field={field} />
+            {/* <CivitaiModelRegistry field={field} /> */}
+            {/* <span>{field.value.length} selected</span> */}
+            {field.value && (
+              <ScrollArea className="w-full bg-gray-100 mx-auto rounded-lg mt-2">
+                <Textarea
+                  className="min-h-[150px] max-h-[300px] p-2 rounded-lg text-xs w-full"
+                  value={JSON.stringify(field.value, null, 2)}
+                  onChange={(e) => {
+                    field.onChange(JSON.parse(e.target.value));
+                  }}
+                />
+              </ScrollArea>
+            )}
+          </div>
+        </AccordionContent>
+      </AccordionItem>
+    </Accordion>
+  );
+}

web/src/components/docs/Navigation.tsx

@@ -242,6 +242,10 @@ export const navigation: Array<NavGroup> = [
     title: "API",
     links: [{ title: "Endpoints", href: "/docs/endpoints" }],
   },
+  {
+    title: "Video Tutorials",
+    links: [{ title: "Archive", href: "/docs/videos" }],
+  },
 ];
 
 export function Navigation(props: React.ComponentPropsWithoutRef<"nav">) {

web/src/components/ui/auto-form/config.ts

@@ -1,3 +1,4 @@
+import AutoFormGPUPicker from "@/components/custom-form/gpu-picker";
 import AutoFormCheckbox from "./fields/checkbox";
 import AutoFormDate from "./fields/date";
 import AutoFormEnum from "./fields/enum";
@@ -8,6 +9,7 @@ import AutoFormSwitch from "./fields/switch";
 import AutoFormTextarea from "./fields/textarea";
 import AutoFormModelsPicker from "@/components/custom-form/model-picker";
 import AutoFormSnapshotPicker from "@/components/custom-form/snapshot-picker";
+import AutoFormModelsPickerUrl from "@/components/custom-form/model-picker-url-only";
 
 export const INPUT_COMPONENTS = {
   checkbox: AutoFormCheckbox,
@@ -22,6 +24,8 @@ export const INPUT_COMPONENTS = {
   // Customs
   snapshot: AutoFormSnapshotPicker,
   models: AutoFormModelsPicker,
+  gpuPicker: AutoFormGPUPicker,
+  modelUrlPicker: AutoFormModelsPickerUrl,
 };
 
 /**

web/src/components/ui/auto-form/types.ts

@@ -1,3 +1,4 @@
+import type { getCurrentPlanWithAuth } from "@/server/getCurrentPlan";
 import type { INPUT_COMPONENTS } from "./config";
 import type { ControllerRenderProps, FieldValues } from "react-hook-form";
 import type * as z from "zod";
@@ -6,6 +7,7 @@ export type FieldConfigItem = {
   description?: React.ReactNode;
   inputProps?: React.InputHTMLAttributes<HTMLInputElement> & {
     showLabel?: boolean;
+    sub?: Awaited<ReturnType<typeof getCurrentPlanWithAuth>>;
   };
   fieldType?:
     | keyof typeof INPUT_COMPONENTS

web/src/components/useCurrentPlan.tsx

@@ -0,0 +1,36 @@
+"use client";
+
+import { getCurrentPlanWithAuth } from "@/server/getCurrentPlan";
+import * as React from "react";
+import { createContext, useContext } from "react";
+
+type CurrentPlanContextType = Awaited<
+  ReturnType<typeof getCurrentPlanWithAuth>
+>;
+const CurrentPlanContext = createContext<CurrentPlanContextType | undefined>(
+  undefined,
+);
+
+export function SubscriptionProvider({
+  sub,
+  children,
+}: {
+  sub: CurrentPlanContextType;
+  children: React.ReactNode;
+}) {
+  return (
+    <CurrentPlanContext.Provider value={sub}>
+      {children}
+    </CurrentPlanContext.Provider>
+  );
+}
+
+export const useCurrentPlan = (): CurrentPlanContextType => {
+  const context = useContext(CurrentPlanContext);
+
+  // if (context === undefined) {
+  //   throw new Error("useCurrentPlan must be used within a CurrentPlanProvider");
+  // }
+
+  return context;
+};

web/src/db/schema.ts

@@ -380,12 +380,18 @@ export const resourceUpload = pgEnum("resource_upload", [
 
 export const modelUploadType = pgEnum("model_upload_type", [
   "civitai",
-  "huggingface",
+  "download-url",
+  "huggingface", // remove?
   "other",
 ]);
 
 // https://www.answeroverflow.com/m/1125106227387584552 
-const modelTypes = ["checkpoint", "lora", "embedding", "vae"] as const;
+export const modelTypes = [
+  "checkpoint",
+  "lora",
+  "embedding",
+  "vae",
+] as const
 export const modelType = pgEnum("model_type", modelTypes);
 export type modelEnumType = (typeof modelTypes)[number];
 
@@ -399,8 +405,7 @@ export const modelTable = dbSchema.table("models", {
     .notNull()
     .references(() => userVolume.id, {
       onDelete: "cascade",
-    })
+    }),
-    .notNull(),
 
   model_name: text("model_name"),
   folder_path: text("folder_path"), // in volume
@@ -413,8 +418,10 @@ export const modelTable = dbSchema.table("models", {
     z.infer<typeof CivitaiModelResponse>
   >(),
 
+  // for our own storage
   hf_url: text("hf_url"),
   s3_url: text("s3_url"),
+
   user_url: text("client_url"),
 
   is_public: boolean("is_public").notNull().default(true),
@@ -484,16 +491,6 @@ export const subscriptionStatusTable = dbSchema.table("subscription_status", {
   updated_at: timestamp("updated_at").defaultNow().notNull(),
 });
 
-export const insertCivitaiModelSchema = createInsertSchema(modelTable, {
-  civitai_url: (schema) =>
-    schema.civitai_url
-      .trim()
-      .url({ message: "URL required" })
-      .includes("civitai.com/models", {
-        message: "civitai.com/models link required",
-      }),
-});
-
 export type UserType = InferSelectModel<typeof usersTable>;
 export type WorkflowType = InferSelectModel<typeof workflowTable>;
 export type MachineType = InferSelectModel<typeof machinesTable>;

web/src/server/addCivitaiModelSchema.tsx

@@ -1,5 +1,8 @@
-import { insertCivitaiModelSchema } from "@/db/schema";
+import { z } from "zod";
+import { modelTypes } from "@/db/schema";
 
-export const addCivitaiModelSchema = insertCivitaiModelSchema.pick({
+export const downloadUrlModelSchema = z.object({
-  civitai_url: true,
+  url: z.string().url(),
+  model_type: z.enum(modelTypes).default("checkpoint")
 });
+

web/src/server/addMachineSchema.ts

@@ -17,7 +17,7 @@ export const insertCustomMachineSchema = createInsertSchema(machinesTable, {
       comfyui: "d0165d819afe76bd4e6bdd710eb5f3e571b6a804",
       git_custom_nodes: {
         "https://github.com/BennyKok/comfyui-deploy.git": {
-          hash: "43fe0a384aa5fa9e141d4a264b2ed40a73b817bc",
+          hash: "a838cb7ad425e5652c3931fbafdc886b53c48a22",
           disabled: false,
         },
       },

web/src/server/curdModel.ts

@@ -2,6 +2,7 @@
 
 import { auth } from "@clerk/nextjs";
 import {
+  modelEnumType,
   modelTable,
   ModelType,
   userVolume,
@@ -11,7 +12,7 @@ import { withServerPromise } from "./withServerPromise";
 import { db } from "@/db/db";
 import type { z } from "zod";
 import { headers } from "next/headers";
-import { addCivitaiModelSchema } from "./addCivitaiModelSchema";
+import { downloadUrlModelSchema } from "./addCivitaiModelSchema";
 import { and, eq, isNull } from "drizzle-orm";
 import { CivitaiModelResponse, getModelTypeDetails } from "@/types/civitai";
 
@@ -109,14 +110,150 @@ function getUrl(civitai_url: string) {
   return { url: baseUrl + modelId, modelVersionId };
 }
 
+// Helper function to make a HEAD request and follow redirects
+async function fetchFinalUrl(
+  url: string,
+): Promise<{ finalUrl: string; dispositionFilename?: string }> {
+  console.log("fetching");
+  const response = await fetch(url, { method: "HEAD", redirect: "follow" });
+  if (!response.ok) {
+    console.log("response not ok");
+    throw new Error(`Request failed with status ${response.status}`);
+  }
+  const contentDisposition = response.headers.get("content-disposition");
+  let filename;
+  if (contentDisposition) {
+    const matches = contentDisposition.match(
+      /filename\*?=['"]?(?:UTF-8'')?([^;'"\n]*)['"]?;?/i,
+    );
+    filename = matches && matches[1]
+      ? decodeURIComponent(matches[1])
+      : undefined;
+  }
+  return { finalUrl: response.url, dispositionFilename: filename };
+}
+
+// The main function for validation
+export const addModel = withServerPromise(
+  async (data: z.infer<typeof downloadUrlModelSchema>) => {
+    const { url } = data;
+
+    if (url.includes("civitai.com/models/")) {
+      // Make a HEAD request to check for 200 OK
+      const response = await fetch(url, { method: "HEAD" });
+      if (!response.ok) {
+        createModelErrorRecord(
+          url,
+          `civitai gave non-ok response`,
+          "civitai",
+          data.model_type,
+        );
+      }
+      addCivitaiModel(data);
+    } else {
+      const { finalUrl, dispositionFilename } = await fetchFinalUrl(url);
+      console.log("finished fetching");
+      console.log(finalUrl, dispositionFilename);
+
+      if (!dispositionFilename) {
+        console.log("no file name");
+        createModelErrorRecord(
+          url,
+          `Could not find a filename from resolved Url: ${finalUrl}`,
+          "download-url",
+          data.model_type,
+        );
+        return;
+      }
+
+      const validExtensions = [".ckpt", ".pt", ".bin", ".pth", ".safetensors"];
+      const extension = dispositionFilename.slice(
+        dispositionFilename.lastIndexOf("."),
+      );
+      if (!validExtensions.includes(extension)) {
+        console.log("invalid extension");
+        createModelErrorRecord(
+          url,
+          `file ext ${extension} is invalid. Valid extensions: ${validExtensions}`,
+          "download-url",
+          data.model_type,
+        );
+      }
+      addModelDownloadUrl(data, dispositionFilename);
+    }
+  },
+);
+
+export const addModelDownloadUrl = withServerPromise(
+  async (data: z.infer<typeof downloadUrlModelSchema>, filename: string) => {
+    console.log("adding model download");
+    const { userId, orgId } = auth();
+    if (!userId) return { error: "No user id" };
+    const volumes = await retrieveModelVolumes();
+
+    const a = await db
+      .insert(modelTable)
+      .values({
+        user_id: userId,
+        org_id: orgId,
+        upload_type: "download-url",
+        model_name: filename,
+        user_url: data.url,
+        user_volume_id: volumes[0].id,
+        model_type: data.model_type,
+      })
+      .returning();
+
+    const b = a[0];
+    console.log("download url about to upload");
+    await uploadModel(data, b, volumes[0]);
+  },
+);
+
+export const getCivitaiModelRes = async (civitaiUrl: string) => {
+  const { url, modelVersionId } = getUrl(civitaiUrl);
+  const civitaiModelRes = await fetch(url)
+    .then((x) => x.json())
+    .then((a) => {
+      return CivitaiModelResponse.parse(a);
+    });
+  return { civitaiModelRes, url, modelVersionId };
+};
+
+const createModelErrorRecord = async (
+  url: string,
+  errorMessage: string,
+  upload_type: "civitai" | "download-url",
+  model_type: modelEnumType,
+) => {
+  const { userId, orgId } = auth();
+  if (!userId) return { error: "No user id" };
+  const volumes = await retrieveModelVolumes();
+
+  const a = await db
+    .insert(modelTable)
+    .values({
+      user_id: userId,
+      org_id: orgId,
+      user_volume_id: volumes[0].id,
+      upload_type: "civitai",
+      model_type,
+      civitai_url: upload_type === "civitai" ? url : undefined,
+      user_url: upload_type === "download-url" ? url : undefined,
+      error_log: errorMessage,
+      status: "failed",
+    })
+    .returning();
+  return a;
+};
+
 export const addCivitaiModel = withServerPromise(
-  async (data: z.infer<typeof addCivitaiModelSchema>) => {
+  async (data: z.infer<typeof downloadUrlModelSchema>) => {
     const { userId, orgId } = auth();
 
-    if (!data.civitai_url) return { error: "no civitai_url" };
     if (!userId) return { error: "No user id" };
 
-    const { url, modelVersionId } = getUrl(data?.civitai_url);
+    const { url, modelVersionId } = getUrl(data.url);
     const civitaiModelRes = await fetch(url)
       .then((x) => x.json())
       .then((a) => {
@@ -142,18 +279,17 @@ export const addCivitaiModel = withServerPromise(
       selectedModelVersionId = selectedModelVersion?.id.toString();
     }
 
-    const userVolume = await getModelVolumes();
+    const volumes = await retrieveModelVolumes();
-    let cVolume;
-    if (userVolume.length === 0) {
-      const volume = await addModelVolume();
-      cVolume = volume[0];
-    } else {
-      cVolume = userVolume[0];
-    }
 
     const model_type = getModelTypeDetails(civitaiModelRes.type);
     if (!model_type) {
-      return 
+      createModelErrorRecord(
+        url,
+        `Civitai model type ${civitaiModelRes.type} is not currently supported`,
+        "civitai",
+        data.model_type,
+      );
+      return;
     }
 
     const a = await db
@@ -165,18 +301,17 @@ export const addCivitaiModel = withServerPromise(
         model_name: selectedModelVersion.files[0].name,
         civitai_id: civitaiModelRes.id.toString(),
         civitai_version_id: selectedModelVersionId,
-        civitai_url: data.civitai_url,
+        civitai_url: data.url, // TODO: need to confirm
         civitai_download_url: selectedModelVersion.files[0].downloadUrl,
         civitai_model_response: civitaiModelRes,
-        user_volume_id: cVolume.id,
+        user_volume_id: volumes[0].id,
         model_type,
-        updated_at: new Date(),
       })
       .returning();
 
     const b = a[0];
 
-    await uploadModel(data, b, cVolume);
+    await uploadModel(data, b, volumes[0]);
   },
 );
 
@@ -216,7 +351,7 @@ export const addCivitaiModel = withServerPromise(
 // );
 
 async function uploadModel(
-  data: z.infer<typeof addCivitaiModelSchema>,
+  data: z.infer<typeof downloadUrlModelSchema>,
   c: ModelType,
   v: UserVolumeType,
 ) {
@@ -238,7 +373,9 @@ async function uploadModel(
         "Content-Type": "application/json",
       },
       body: JSON.stringify({
-        download_url: c.civitai_download_url,
+        download_url: c.upload_type === "civitai"
+          ? c.civitai_download_url
+          : c.user_url,
         volume_name: v.volume_name,
         volume_id: v.id,
         model_id: c.id,
@@ -253,7 +390,6 @@ async function uploadModel(
     await db
       .update(modelTable)
       .set({
-        ...data,
         status: "failed",
         error_log: error_log,
       })

web/src/server/getCurrentPlan.tsx

@@ -3,6 +3,7 @@ import { and, desc, eq, isNull, or } from "drizzle-orm";
 import { subscriptionStatusTable } from "@/db/schema";
 import { APIKeyUserType } from "@/server/APIKeyBodyRequest";
 import { auth } from "@clerk/nextjs";
+import "server-only";
 
 export async function getCurrentPlanWithAuth() {
   const { userId, orgId } = auth();
@@ -23,7 +24,10 @@ export async function getCurrentPlan({ user_id, org_id }: APIKeyUserType) {
       eq(subscriptionStatusTable.user_id, user_id),
       org_id
         ? eq(subscriptionStatusTable.org_id, org_id)
-        : or(isNull(subscriptionStatusTable.org_id), eq(subscriptionStatusTable.org_id, "")),
+        : or(
+            isNull(subscriptionStatusTable.org_id),
+            eq(subscriptionStatusTable.org_id, ""),
+          ),
     ),
     orderBy: desc(subscriptionStatusTable.created_at),
   });