inference-server 1.0.0-beta.19

package/src/engines/stable-diffusion-cpp/validateModelFiles.ts

@@ -0,0 +1,119 @@
+import fs from 'node:fs'
+import { calculateFileChecksum } from '#package/lib/calculateFileChecksum.js'
+import { resolveModelFileLocation } from '#package/lib/resolveModelFileLocation.js'
+import { ModelFileSource } from '#package/types/index.js'
+import { StableDiffusionModelConfig } from './engine.js'
+
+interface ModelValidationErrors {
+	model?: string
+	clipL?: string
+	clipG?: string
+	vae?: string
+	t5xxl?: string
+	controlNet?: string
+	taesd?: string
+	lora?: {
+		[index: number]: string
+	}
+}
+
+export interface ModelValidationResult {
+	message: string
+	errors: ModelValidationErrors
+}
+
+export async function validateModelFiles(
+	config: StableDiffusionModelConfig,
+): Promise<ModelValidationResult | undefined> {
+	const validateFile = async (component: string, src: ModelFileSource) => {
+		const fileLocation = resolveModelFileLocation({
+			url: src.url,
+			filePath: src.file,
+			modelsCachePath: config.modelsCachePath,
+		})
+		if (!fs.existsSync(fileLocation)) {
+			return {
+				component,
+				message: `${component} file missing at ${fileLocation}`,
+			}
+		}
+		const ipullFile = fileLocation + '.ipull'
+		let validatedChecksum = false
+		if (fs.existsSync(ipullFile)) {
+			// if we have a valid file at the download destination, we can remove the ipull file
+			if (src.sha256) {
+				const fileHash = await calculateFileChecksum(fileLocation, 'sha256')
+				if (fileHash === src.sha256) {
+					fs.unlinkSync(ipullFile)
+					validatedChecksum = true
+				}
+			}
+			if (!validatedChecksum) {
+				return {
+					component,
+					message: `${component} with incomplete download`,
+				}
+			}
+		}
+
+		if (!validatedChecksum && src.sha256) {
+			const fileHash = await calculateFileChecksum(fileLocation, 'sha256')
+			if (fileHash !== src.sha256) {
+				return {
+					component,
+					message: `${component} file sha256 checksum mismatch: expected ${src.sha256} got ${fileHash} at ${fileLocation}`,
+				}
+			}
+		}
+		return undefined
+	}
+
+	const validationPromises = []
+	if (config.clipL) {
+		validationPromises.push(validateFile('clipL', config.clipL))
+	}
+	if (config.clipG) {
+		validationPromises.push(validateFile('clipG', config.clipG))
+	}
+	if (config.vae) {
+		validationPromises.push(validateFile('vae', config.vae))
+	}
+	if (config.t5xxl) {
+		validationPromises.push(validateFile('t5xxl', config.t5xxl))
+	}
+	if (config.controlNet) {
+		validationPromises.push(validateFile('controlNet', config.controlNet))
+	}
+	if (config.taesd) {
+		validationPromises.push(validateFile('taesd', config.taesd))
+	}
+	if (config.sha256) {
+		validationPromises.push(validateFile('model', { file: config.location, sha256: config.sha256 }))
+	}
+
+	// const loraDir = path.join(path.dirname(config.location), 'lora')
+	if (config.loras) {
+		for (const lora of config.loras) {
+			validationPromises.push(validateFile('lora', lora))
+
+			// const loraFile = path.join(loraDir, getModelFileName(lora))
+			// if (!fs.existsSync(loraFile)) {
+			// 	return `lora file missing: ${loraFile}`
+			// }
+		}
+	}
+
+	const res = await Promise.all(validationPromises)
+	const validationErrors = res.filter((e) => !!e)
+	if (validationErrors.length) {
+		return {
+			message: 'Invalid model files',
+			errors: validationErrors.reduce((acc, e) => {
+				acc[e.component as keyof ModelValidationErrors] = e.message
+				return acc
+			}, {} as ModelValidationErrors),
+		}
+	}
+
+	return undefined
+}