From 3e8ef33cbc22ad94f7144de87918a1e34a02b4cc Mon Sep 17 00:00:00 2001
From: Christian Byrne <cbyrne@comfy.org>
Date: Wed, 12 Mar 2025 07:43:15 -0700
Subject: [PATCH] Fetch model metadata for Civitai models embedded in workflows
 (#2994)

---
 src/composables/useCivitaiModel.ts | 95 ++++++++++++++++++++++++++++++
 src/composables/useDownload.ts     | 23 ++++++--
 src/utils/formatUtil.ts            | 22 +++++++
 3 files changed, 136 insertions(+), 4 deletions(-)
 create mode 100644 src/composables/useCivitaiModel.ts

diff --git a/src/composables/useCivitaiModel.ts b/src/composables/useCivitaiModel.ts
new file mode 100644
index 000000000..c47b27483
--- /dev/null
+++ b/src/composables/useCivitaiModel.ts
@@ -0,0 +1,95 @@
+import { useAsyncState } from '@vueuse/core'
+import { computed } from 'vue'
+
+type ModelType =
+  | 'Checkpoint'
+  | 'TextualInversion'
+  | 'Hypernetwork'
+  | 'AestheticGradient'
+  | 'LORA'
+  | 'Controlnet'
+  | 'Poses'
+
+interface CivitaiFileMetadata {
+  fp?: 'fp16' | 'fp32'
+  size?: 'full' | 'pruned'
+  format?: 'SafeTensor' | 'PickleTensor' | 'Other'
+}
+
+interface CivitaiModelFile {
+  name: string
+  id: number
+  sizeKB: number
+  type: string
+  downloadUrl: string
+  metadata: CivitaiFileMetadata
+}
+
+interface CivitaiModel {
+  name: string
+  type: ModelType
+}
+
+interface CivitaiModelVersionResponse {
+  id: number
+  name: string
+  model: CivitaiModel
+  modelId: number
+  files: CivitaiModelFile[]
+  [key: string]: any
+}
+
+/**
+ * Composable to manage Civitai model
+ * @param url - The URL of the Civitai model, where the model ID is the last part of the URL's pathname
+ * @see https://developer.civitai.com/docs/api/public-rest
+ * @example
+ * const { fileSize, isLoading, error, modelData } =
+ *  useCivitaiModel('https://civitai.com/api/download/models/16576?type=Model&format=SafeTensor&size=full&fp=fp16')
+ */
+export function useCivitaiModel(url: string) {
+  const createModelVersionUrl = (modelId: string): string =>
+    `https://civitai.com/api/v1/model-versions/${modelId}`
+
+  const extractModelIdFromUrl = (): string | null => {
+    const urlObj = new URL(url)
+    return urlObj.pathname.split('/').pop() || null
+  }
+
+  const fetchModelData =
+    async (): Promise<CivitaiModelVersionResponse | null> => {
+      const modelId = extractModelIdFromUrl()
+      if (!modelId) return null
+
+      const apiUrl = createModelVersionUrl(modelId)
+      const res = await fetch(apiUrl)
+      return res.json()
+    }
+
+  const findMatchingFileSize = (): number | null => {
+    const matchingFile = modelData.value?.files?.find(
+      (file) => file.downloadUrl && url.startsWith(file.downloadUrl)
+    )
+
+    return matchingFile?.sizeKB ? matchingFile.sizeKB << 10 : null
+  }
+
+  const {
+    state: modelData,
+    isLoading,
+    error
+  } = useAsyncState(fetchModelData, null, {
+    immediate: true
+  })
+
+  const fileSize = computed(() =>
+    !isLoading.value ? findMatchingFileSize() : null
+  )
+
+  return {
+    fileSize,
+    isLoading,
+    error,
+    modelData
+  }
+}
diff --git a/src/composables/useDownload.ts b/src/composables/useDownload.ts
index acbbf863e..22d9d3a42 100644
--- a/src/composables/useDownload.ts
+++ b/src/composables/useDownload.ts
@@ -1,16 +1,24 @@
+import { whenever } from '@vueuse/core'
 import { onMounted, ref } from 'vue'
 
+import { useCivitaiModel } from '@/composables/useCivitaiModel'
+import { isCivitaiModelUrl } from '@/utils/formatUtil'
+
 export function useDownload(url: string, fileName?: string) {
   const fileSize = ref<number | null>(null)
 
-  const fetchFileSize = async (): Promise<number | null> => {
+  const setFileSize = (size: number) => {
+    fileSize.value = size
+  }
+
+  const fetchFileSize = async () => {
     try {
       const response = await fetch(url, { method: 'HEAD' })
       if (!response.ok) throw new Error('Failed to fetch file size')
 
       const size = response.headers.get('content-length')
       if (size) {
-        return parseInt(size)
+        setFileSize(parseInt(size))
       } else {
         console.error('"content-length" header not found')
         return null
@@ -33,8 +41,15 @@ export function useDownload(url: string, fileName?: string) {
     link.click()
   }
 
-  onMounted(async () => {
-    fileSize.value = await fetchFileSize()
+  onMounted(() => {
+    if (isCivitaiModelUrl(url)) {
+      const { fileSize: civitaiSize, error: civitaiErr } = useCivitaiModel(url)
+      whenever(civitaiSize, setFileSize)
+      // Try falling back to normal fetch if using Civitai API fails
+      whenever(civitaiErr, fetchFileSize, { once: true })
+    } else {
+      fetchFileSize()
+    }
   })
 
   return {
diff --git a/src/utils/formatUtil.ts b/src/utils/formatUtil.ts
index 659b96c9e..05390c668 100644
--- a/src/utils/formatUtil.ts
+++ b/src/utils/formatUtil.ts
@@ -343,3 +343,25 @@ export const generateUUID = (): string => {
  */
 export const formatNumber = (num?: number): string =>
   num?.toLocaleString() ?? 'N/A'
+
+/**
+ * Checks if a URL is a Civitai model URL
+ * @example
+ * isCivitaiModelUrl('https://civitai.com/api/download/models/1234567890') // true
+ * isCivitaiModelUrl('https://civitai.com/api/v1/models/1234567890') // true
+ * isCivitaiModelUrl('https://civitai.com/api/v1/models-versions/15342') // true
+ * isCivitaiModelUrl('https://example.com/model.safetensors') // false
+ */
+export const isCivitaiModelUrl = (url: string): boolean => {
+  if (!isValidUrl(url)) return false
+  if (!url.includes('civitai.com')) return false
+
+  const urlObj = new URL(url)
+  const pathname = urlObj.pathname
+
+  return (
+    /^\/api\/download\/models\/(\d+)$/.test(pathname) ||
+    /^\/api\/v1\/models\/(\d+)$/.test(pathname) ||
+    /^\/api\/v1\/models-versions\/(\d+)$/.test(pathname)
+  )
+}