From 6fe2297cc1b8c7c6aa7fd6671c8f09561172987c Mon Sep 17 00:00:00 2001
From: Robin Huang <robin.j.huang@gmail.com>
Date: Tue, 6 Aug 2024 10:45:07 -0700
Subject: [PATCH] Add models information to default workflow. (#321)

* Add models information to default workflow.

* Add models to zod schema.

* Fix zod schema.

* Update schema name.

* Add z prefix to modelfile schema name.
---
 src/scripts/defaultGraph.ts | 11 ++++++++++-
 src/types/comfyWorkflow.ts  | 12 +++++++++++-
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/src/scripts/defaultGraph.ts b/src/scripts/defaultGraph.ts
index ebf7c5b10..753a18ce7 100644
--- a/src/scripts/defaultGraph.ts
+++ b/src/scripts/defaultGraph.ts
@@ -133,5 +133,14 @@ export const defaultGraph: ComfyWorkflowJSON = {
   groups: [],
   config: {},
   extra: {},
-  version: 0.4
+  version: 0.4,
+  models: [
+    {
+      name: 'v1-5-pruned-emaonly.ckpt',
+      url: 'https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/v1-5-pruned-emaonly.ckpt',
+      hash: 'cc6cb27103417325ff94f52b7a5d2dde45a7515b25c255d8e396c90014281516',
+      hash_type: 'SHA256',
+      directory: 'checkpoints'
+    }
+  ]
 }
diff --git a/src/types/comfyWorkflow.ts b/src/types/comfyWorkflow.ts
index bdc5eebe5..071bdce5f 100644
--- a/src/types/comfyWorkflow.ts
+++ b/src/types/comfyWorkflow.ts
@@ -1,6 +1,15 @@
 import { z } from 'zod'
 import { fromZodError } from 'zod-validation-error'
 
+// Definition of an AI model file used in the workflow.
+const zModelFile = z.object({
+  name: z.string(),
+  url: z.string().url(),
+  hash: z.string(),
+  hash_type: z.string(),
+  directory: z.string()
+})
+
 const zComfyLink = z.tuple([
   z.number(), // Link id
   z.number(), // Node id of source node
@@ -121,7 +130,8 @@ export const zComfyWorkflow = z
     groups: z.array(zGroup).optional(),
     config: zConfig.optional().nullable(),
     extra: zExtra.optional().nullable(),
-    version: z.number()
+    version: z.number(),
+    models: z.array(zModelFile).optional()
   })
   .passthrough()