- table.getColumn("name")?.setFilterValue(event.target.value)
- }
+ table.getColumn("model_name")?.setFilterValue(event.target.value)}
className="max-w-sm"
/>
@@ -268,17 +302,17 @@ export function CheckpointList({ data }: { data: CheckpointItemList[] }) {
// TODO: limitations based on plan
}
tooltip={"Add models using their civitai url!"}
- title="Civitai Checkpoint"
+ title="Add a Civitai Model"
description="Pick a model from civitai"
- serverAction={addCivitaiCheckpoint}
- formSchema={addCivitaiCheckpointSchema}
+ serverAction={addCivitaiModel}
+ formSchema={addCivitaiModelSchema}
fieldConfig={{
civitai_url: {
fieldType: "fallback",
inputProps: { required: true },
description: (
<>
- Pick a checkpoint from{" "}
+ Pick a model from{" "}
{
return (
- {header.isPlaceholder
- ? null
- : flexRender(
- header.column.columnDef.header,
- header.getContext()
- )}
+ {header.isPlaceholder ? null : flexRender(
+ header.column.columnDef.header,
+ header.getContext(),
+ )}
);
})}
@@ -315,32 +347,34 @@ export function CheckpointList({ data }: { data: CheckpointItemList[] }) {
))}
- {table.getRowModel().rows?.length ? (
- table.getRowModel().rows.map((row) => (
-
- {row.getVisibleCells().map((cell) => (
-
- {flexRender(
- cell.column.columnDef.cell,
- cell.getContext()
- )}
-
- ))}
+ {table.getRowModel().rows?.length
+ ? (
+ table.getRowModel().rows.map((row) => (
+
+ {row.getVisibleCells().map((cell) => (
+
+ {flexRender(
+ cell.column.columnDef.cell,
+ cell.getContext(),
+ )}
+
+ ))}
+
+ ))
+ )
+ : (
+
+
+ No results.
+
- ))
- ) : (
-
-
- No results.
-
-
- )}
+ )}
diff --git a/web/src/components/custom-form/checkpoint-input.tsx b/web/src/components/custom-form/model-input.tsx
similarity index 95%
rename from web/src/components/custom-form/checkpoint-input.tsx
rename to web/src/components/custom-form/model-input.tsx
index 9a21cb9..66cf7f5 100644
--- a/web/src/components/custom-form/checkpoint-input.tsx
+++ b/web/src/components/custom-form/model-input.tsx
@@ -7,7 +7,7 @@ import AutoFormInput from "../ui/auto-form/fields/input";
import { useDebouncedCallback } from "use-debounce";
import { CivitaiModelResponse } from "@/types/civitai";
import { z } from "zod";
-import { insertCivitaiCheckpointSchema } from "@/db/schema";
+import { insertCivitaiModelSchema } from "@/db/schema";
function getUrl(civitai_url: string) {
// expect to be a URL to be https://civitai.com/models/36520
@@ -33,7 +33,7 @@ export default function AutoFormCheckpointInput(
const handleSearch = useDebouncedCallback((search) => {
const validationResult =
- insertCivitaiCheckpointSchema.shape.civitai_url.safeParse(search);
+ insertCivitaiModelSchema.shape.civitai_url.safeParse(search);
if (!validationResult.success) {
console.error(validationResult.error);
// Optionally set an error state here
diff --git a/web/src/db/schema.ts b/web/src/db/schema.ts
index da3433c..5a560a3 100644
--- a/web/src/db/schema.ts
+++ b/web/src/db/schema.ts
@@ -12,7 +12,7 @@ import {
real,
} from "drizzle-orm/pg-core";
import { createInsertSchema, createSelectSchema } from "drizzle-zod";
-import { z } from "zod";
+import { TypeOf, z } from "zod";
export const dbSchema = pgSchema("comfyui_deploy");
@@ -376,15 +376,25 @@ export const modelUploadType = pgEnum("model_upload_type", [
"other",
]);
-export const checkpointTable = dbSchema.table("checkpoints", {
+// https://www.answeroverflow.com/m/1125106227387584552
+const modelTypes = [
+ "checkpoint",
+ "lora",
+ "embedding",
+ "vae",
+] as const
+export const modelType = pgEnum("model_type", modelTypes);
+export type modelEnumType = typeof modelTypes[number]
+
+export const modelTable = dbSchema.table("models", {
id: uuid("id").primaryKey().defaultRandom().notNull(),
- user_id: text("user_id").references(() => usersTable.id, {}), // perhaps a "special" user_id for global checkpoints
+ user_id: text("user_id").references(() => usersTable.id, {}), // perhaps a "special" user_id for global models
org_id: text("org_id"),
description: text("description"),
- checkpoint_volume_id: uuid("checkpoint_volume_id")
+ user_volume_id: uuid("user_volume_id")
.notNull()
- .references(() => checkpointVolumeTable.id, {
+ .references(() => userVolume.id, {
onDelete: "cascade",
})
.notNull(),
@@ -408,12 +418,13 @@ export const checkpointTable = dbSchema.table("checkpoints", {
status: resourceUpload("status").notNull().default("started"),
upload_machine_id: text("upload_machine_id"), // TODO: review if actually needed
upload_type: modelUploadType("upload_type").notNull(),
+ model_type: modelType("model_type").notNull(),
error_log: text("error_log"),
created_at: timestamp("created_at").defaultNow().notNull(),
updated_at: timestamp("updated_at").defaultNow().notNull(),
});
-export const checkpointVolumeTable = dbSchema.table("checkpoint_volume", {
+export const userVolume = dbSchema.table("user_volume", {
id: uuid("id").primaryKey().defaultRandom().notNull(),
user_id: text("user_id").references(() => usersTable.id, {
// onDelete: "cascade",
@@ -425,23 +436,23 @@ export const checkpointVolumeTable = dbSchema.table("checkpoint_volume", {
disabled: boolean("disabled").default(false).notNull(),
});
-export const checkpointRelations = relations(checkpointTable, ({ one }) => ({
+export const modelRelations = relations(modelTable, ({ one }) => ({
user: one(usersTable, {
- fields: [checkpointTable.user_id],
+ fields: [modelTable.user_id],
references: [usersTable.id],
}),
- volume: one(checkpointVolumeTable, {
- fields: [checkpointTable.checkpoint_volume_id],
- references: [checkpointVolumeTable.id],
+ volume: one(userVolume, {
+ fields: [modelTable.user_volume_id],
+ references: [userVolume.id],
}),
}));
-export const checkpointVolumeRelations = relations(
- checkpointVolumeTable,
+export const modalVolumeRelations = relations(
+ userVolume,
({ many, one }) => ({
- checkpoint: many(checkpointTable),
+ model: many(modelTable),
user: one(usersTable, {
- fields: [checkpointVolumeTable.user_id],
+ fields: [userVolume.user_id],
references: [usersTable.id],
}),
})
@@ -473,8 +484,8 @@ export const subscriptionStatusTable = dbSchema.table("subscription_status", {
updated_at: timestamp("updated_at").defaultNow().notNull(),
});
-export const insertCivitaiCheckpointSchema = createInsertSchema(
- checkpointTable,
+export const insertCivitaiModelSchema = createInsertSchema(
+ modelTable,
{
civitai_url: (schema) =>
schema.civitai_url
@@ -491,8 +502,8 @@ export type WorkflowType = InferSelectModel;
export type MachineType = InferSelectModel;
export type WorkflowVersionType = InferSelectModel;
export type DeploymentType = InferSelectModel;
-export type CheckpointType = InferSelectModel;
-export type CheckpointVolumeType = InferSelectModel<
- typeof checkpointVolumeTable
+export type ModelType = InferSelectModel;
+export type UserVolumeType = InferSelectModel<
+ typeof userVolume
>;
export type UserUsageType = InferSelectModel;
diff --git a/web/src/server/addCheckpointSchema.tsx b/web/src/server/addCheckpointSchema.tsx
deleted file mode 100644
index b83f5a3..0000000
--- a/web/src/server/addCheckpointSchema.tsx
+++ /dev/null
@@ -1,5 +0,0 @@
-import { insertCivitaiCheckpointSchema } from "@/db/schema";
-
-export const addCivitaiCheckpointSchema = insertCivitaiCheckpointSchema.pick({
- civitai_url: true,
-});
diff --git a/web/src/server/addCivitaiModelSchema.tsx b/web/src/server/addCivitaiModelSchema.tsx
new file mode 100644
index 0000000..3f205eb
--- /dev/null
+++ b/web/src/server/addCivitaiModelSchema.tsx
@@ -0,0 +1,5 @@
+import { insertCivitaiModelSchema } from "@/db/schema";
+
+export const addCivitaiModelSchema = insertCivitaiModelSchema.pick({
+ civitai_url: true,
+});
diff --git a/web/src/server/curdMachine.ts b/web/src/server/curdMachine.ts
index 77fea3b..cbe58b8 100644
--- a/web/src/server/curdMachine.ts
+++ b/web/src/server/curdMachine.ts
@@ -15,7 +15,7 @@ import { headers } from "next/headers";
import { redirect } from "next/navigation";
import "server-only";
import type { z } from "zod";
-import { retrieveCheckpointVolumes } from "./curdCheckpoint";
+import { retrieveModelVolumes } from "./curdModel";
export async function getMachines() {
const { userId, orgId } = auth();
@@ -190,7 +190,7 @@ async function _buildMachine(
throw new Error("No domain");
}
- const volumes = await retrieveCheckpointVolumes();
+ const volumes = await retrieveModelVolumes();
// Call remote builder
const result = await fetch(`${process.env.MODAL_BUILDER_URL!}/create`, {
method: "POST",
@@ -204,7 +204,7 @@ async function _buildMachine(
callback_url: `${protocol}://${domain}/api/machine-built`,
models: data.models, //JSON.parse(data.models as string),
gpu: data.gpu && data.gpu.length > 0 ? data.gpu : "T4",
- checkpoint_volume_name: volumes[0].volume_name,
+ model_volume_name: volumes[0].volume_name,
}),
});
diff --git a/web/src/server/curdCheckpoint.ts b/web/src/server/curdModel.ts
similarity index 71%
rename from web/src/server/curdCheckpoint.ts
rename to web/src/server/curdModel.ts
index 3c01532..cd22d16 100644
--- a/web/src/server/curdCheckpoint.ts
+++ b/web/src/server/curdModel.ts
@@ -2,92 +2,91 @@
import { auth } from "@clerk/nextjs";
import {
- checkpointTable,
- CheckpointType,
- checkpointVolumeTable,
- CheckpointVolumeType,
+ modelTable,
+ ModelType,
+ userVolume,
+ UserVolumeType,
} from "@/db/schema";
import { withServerPromise } from "./withServerPromise";
import { db } from "@/db/db";
import type { z } from "zod";
import { headers } from "next/headers";
-import { addCivitaiCheckpointSchema } from "./addCheckpointSchema";
+import { addCivitaiModelSchema } from "./addCivitaiModelSchema";
import { and, eq, isNull } from "drizzle-orm";
-import { CivitaiModelResponse } from "@/types/civitai";
+import { CivitaiModelResponse, getModelTypeDetails } from "@/types/civitai";
-export async function getCheckpoints() {
+export async function getModel() {
const { userId, orgId } = auth();
if (!userId) throw new Error("No user id");
- const checkpoints = await db
+ const models = await db
.select()
- .from(checkpointTable)
+ .from(modelTable)
.where(
orgId
- ? eq(checkpointTable.org_id, orgId)
+ ? eq(modelTable.org_id, orgId)
// make sure org_id is null
: and(
- eq(checkpointTable.user_id, userId),
- isNull(checkpointTable.org_id),
+ eq(modelTable.user_id, userId),
+ isNull(modelTable.org_id),
),
);
- return checkpoints;
+ return models;
}
-export async function getCheckpointById(id: string) {
+export async function getModelById(id: string) {
const { userId, orgId } = auth();
if (!userId) throw new Error("No user id");
- const checkpoint = await db
+ const model = await db
.select()
- .from(checkpointTable)
+ .from(modelTable)
.where(
and(
- orgId ? eq(checkpointTable.org_id, orgId) : and(
- eq(checkpointTable.user_id, userId),
- isNull(checkpointTable.org_id),
+ orgId ? eq(modelTable.org_id, orgId) : and(
+ eq(modelTable.user_id, userId),
+ isNull(modelTable.org_id),
),
- eq(checkpointTable.id, id),
+ eq(modelTable.id, id),
),
);
- return checkpoint[0];
+ return model[0];
}
-export async function getCheckpointVolumes() {
+export async function getModelVolumes() {
const { userId, orgId } = auth();
if (!userId) throw new Error("No user id");
const volume = await db
.select()
- .from(checkpointVolumeTable)
+ .from(userVolume)
.where(
and(
orgId
- ? eq(checkpointVolumeTable.org_id, orgId)
+ ? eq(userVolume.org_id, orgId)
// make sure org_id is null
: and(
- eq(checkpointVolumeTable.user_id, userId),
- isNull(checkpointVolumeTable.org_id),
+ eq(userVolume.user_id, userId),
+ isNull(userVolume.org_id),
),
- eq(checkpointVolumeTable.disabled, false),
+ eq(userVolume.disabled, false),
),
);
return volume;
}
-export async function retrieveCheckpointVolumes() {
- let volumes = await getCheckpointVolumes();
+export async function retrieveModelVolumes() {
+ let volumes = await getModelVolumes();
if (volumes.length === 0) {
// create volume if not already created
- volumes = await addCheckpointVolume();
+ volumes = await addModelVolume();
}
return volumes;
}
-export async function addCheckpointVolume() {
+export async function addModelVolume() {
const { userId, orgId } = auth();
if (!userId) throw new Error("No user id");
- // Insert the new checkpointVolume into the checkpointVolumeTable
const insertedVolume = await db
- .insert(checkpointVolumeTable)
+ .insert(userVolume)
.values({
user_id: userId,
org_id: orgId,
@@ -111,8 +110,8 @@ function getUrl(civitai_url: string) {
return { url: baseUrl + modelId, modelVersionId };
}
-export const addCivitaiCheckpoint = withServerPromise(
- async (data: z.infer) => {
+export const addCivitaiModel = withServerPromise(
+ async (data: z.infer) => {
const { userId, orgId } = auth();
if (!data.civitai_url) return { error: "no civitai_url" };
@@ -145,17 +144,22 @@ export const addCivitaiCheckpoint = withServerPromise(
selectedModelVersionId = selectedModelVersion?.id.toString();
}
- const checkpointVolumes = await getCheckpointVolumes();
+ const userVolume = await getModelVolumes();
let cVolume;
- if (checkpointVolumes.length === 0) {
- const volume = await addCheckpointVolume();
+ if (userVolume.length === 0) {
+ const volume = await addModelVolume();
cVolume = volume[0];
} else {
- cVolume = checkpointVolumes[0];
+ cVolume = userVolume[0];
+ }
+
+ const model_type = getModelTypeDetails(civitaiModelRes.type);
+ if (!model_type) {
+ return
}
const a = await db
- .insert(checkpointTable)
+ .insert(modelTable)
.values({
user_id: userId,
org_id: orgId,
@@ -166,15 +170,15 @@ export const addCivitaiCheckpoint = withServerPromise(
civitai_url: data.civitai_url,
civitai_download_url: selectedModelVersion.files[0].downloadUrl,
civitai_model_response: civitaiModelRes,
- checkpoint_volume_id: cVolume.id,
+ user_volume_id: cVolume.id,
+ model_type,
updated_at: new Date(),
})
.returning();
const b = a[0];
- await uploadCheckpoint(data, b, cVolume);
- // redirect(`/checkpoints/${b.id}`);
+ await uploadModel(data, b, cVolume);
},
);
@@ -213,10 +217,10 @@ export const addCivitaiCheckpoint = withServerPromise(
// },
// );
-async function uploadCheckpoint(
- data: z.infer,
- c: CheckpointType,
- v: CheckpointVolumeType,
+async function uploadModel(
+ data: z.infer,
+ c: ModelType,
+ v: UserVolumeType,
) {
const headersList = headers();
@@ -239,9 +243,9 @@ async function uploadCheckpoint(
download_url: c.civitai_download_url,
volume_name: v.volume_name,
volume_id: v.id,
- checkpoint_id: c.id,
+ model_id: c.id,
callback_url: `${protocol}://${domain}/api/volume-upload`,
- upload_type: "checkpoint"
+ upload_type: c.model_type,
}),
},
);
@@ -249,23 +253,23 @@ async function uploadCheckpoint(
if (!result.ok) {
const error_log = await result.text();
await db
- .update(checkpointTable)
+ .update(modelTable)
.set({
...data,
status: "failed",
error_log: error_log,
})
- .where(eq(checkpointTable.id, c.id));
+ .where(eq(modelTable.id, c.id));
throw new Error(`Error: ${result.statusText} ${error_log}`);
} else {
// setting the build machine id
const json = await result.json();
await db
- .update(checkpointTable)
+ .update(modelTable)
.set({
...data,
upload_machine_id: json.build_machine_instance_id,
})
- .where(eq(checkpointTable.id, c.id));
+ .where(eq(modelTable.id, c.id));
}
}
diff --git a/web/src/server/getAllUserCheckpoints.tsx b/web/src/server/getAllUserModel.tsx
similarity index 64%
rename from web/src/server/getAllUserCheckpoints.tsx
rename to web/src/server/getAllUserModel.tsx
index a12a668..192d968 100644
--- a/web/src/server/getAllUserCheckpoints.tsx
+++ b/web/src/server/getAllUserModel.tsx
@@ -1,18 +1,18 @@
import { db } from "@/db/db";
import {
- checkpointTable,
+ modelTable,
} from "@/db/schema";
import { auth } from "@clerk/nextjs";
import { and, desc, eq, isNull } from "drizzle-orm";
-export async function getAllUserCheckpoints() {
+export async function getAllUserModels() {
const { userId, orgId } = await auth();
if (!userId) {
return null;
}
- const checkpoints = await db.query.checkpointTable.findMany({
+ const models = await db.query.modelTable.findMany({
with: {
user: {
columns: {
@@ -28,14 +28,15 @@ export async function getAllUserCheckpoints() {
civitai_model_response: true,
is_public: true,
upload_type: true,
+ model_type: true,
status: true,
},
- orderBy: desc(checkpointTable.updated_at),
+ orderBy: desc(modelTable.updated_at),
where:
orgId != undefined
- ? eq(checkpointTable.org_id, orgId)
- : and(eq(checkpointTable.user_id, userId), isNull(checkpointTable.org_id)),
+ ? eq(modelTable.org_id, orgId)
+ : and(eq(modelTable.user_id, userId), isNull(modelTable.org_id)),
});
- return checkpoints;
+ return models;
}
diff --git a/web/src/types/civitai.ts b/web/src/types/civitai.ts
index 5e57ed0..52b54b2 100644
--- a/web/src/types/civitai.ts
+++ b/web/src/types/civitai.ts
@@ -1,4 +1,5 @@
-import { z } from "zod";
+import { TypeOf, z } from "zod";
+import { modelEnumType } from "@/db/schema";
// from chatgpt https://chat.openai.com/share/4985d20b-30b1-4a28-87f6-6ebf84a1040e
@@ -110,12 +111,21 @@ export const statsSchema = z.object({
tippedAmountCount: z.number(),
});
+const civitaiModelType = z.enum([
+ "Checkpoint",
+ "TextualInversion",
+ "Hypernetwork",
+ "AestheticGradient",
+ "LORA",
+ "Controlnet",
+ "Poses",
+]);
+
export const CivitaiModelResponse = z.object({
id: z.number(),
name: z.string().nullish(),
description: z.string().nullish(),
- // type: z.enum(["Checkpoint", "Lora"]), // TODO: this will be important to know
- type: z.string(),
+ type: civitaiModelType,
poi: z.boolean().nullish(),
nsfw: z.boolean().nullish(),
allowNoCredit: z.boolean().nullish(),
@@ -127,3 +137,22 @@ export const CivitaiModelResponse = z.object({
tags: z.array(z.string()).nullish(),
modelVersions: z.array(modelVersionSchema),
});
+
+export function getModelTypeDetails(
+ modelType: typeof civitaiModelType["_type"],
+): modelEnumType | undefined {
+ switch (modelType) {
+ case "Checkpoint":
+ return "checkpoint"
+ case "TextualInversion":
+ return "embedding"
+ case "LORA":
+ return "lora"
+ case "AestheticGradient":
+ case "Hypernetwork":
+ case "Controlnet":
+ case "Poses":
+ default:
+ return undefined;
+ }
+}