inference-server 1.0.0-beta.19

package/src/engines/node-llama-cpp/types.ts

@@ -0,0 +1,17 @@
+import {
+	ChatHistoryItem,
+	ChatModelFunctions,
+	LlamaChatResponse,
+	LlamaChatResponseFunctionCall,
+} from 'node-llama-cpp'
+
+export interface LlamaChatResult<T extends ChatModelFunctions = any> {
+	responseText: string | null
+	functionCalls?: LlamaChatResponseFunctionCall<T>[]
+	stopReason: LlamaChatResponse['metadata']['stopReason']
+}
+
+export type ContextShiftStrategy = ((options: {
+	chatHistory: ChatHistoryItem[]
+	metadata: any
+}) => { chatHistory: ChatHistoryItem[]; metadata: any }) | null

package/src/engines/node-llama-cpp/util.ts

@@ -0,0 +1,126 @@
+import fs from 'node:fs'
+import path from 'node:path'
+import {
+	ChatHistoryItem,
+	ChatModelResponse,
+	LlamaTextJSON,
+	ChatModelFunctionCall,
+} from 'node-llama-cpp'
+import { CompletionFinishReason, ChatMessage } from '#package/types/index.js'
+import { flattenMessageTextContent } from '#package/lib/flattenMessageTextContent.js'
+import { LlamaChatResult } from './types.js'
+
+export function mapFinishReason(
+	nodeLlamaCppFinishReason: LlamaChatResult['stopReason'],
+): CompletionFinishReason {
+	switch (nodeLlamaCppFinishReason) {
+		case 'functionCalls':
+			return 'toolCalls'
+		case 'stopGenerationTrigger':
+			return 'stopTrigger'
+		case 'customStopTrigger':
+			return 'stopTrigger'
+		default:
+			return nodeLlamaCppFinishReason
+	}
+}
+
+export function addFunctionCallToChatHistory({
+	chatHistory,
+	functionName,
+	functionDescription,
+	callParams,
+	callResult,
+	rawCall,
+	startsNewChunk
+}: {
+	chatHistory: ChatHistoryItem[],
+	functionName: string,
+	functionDescription?: string,
+	callParams: any,
+	callResult: any,
+	rawCall?: LlamaTextJSON,
+	startsNewChunk?: boolean
+}) {
+	const newChatHistory = chatHistory.slice();
+	if (newChatHistory.length === 0 || newChatHistory[newChatHistory.length - 1]!.type !== "model")
+			newChatHistory.push({
+					type: "model",
+					response: []
+			});
+
+	const lastModelResponseItem = newChatHistory[newChatHistory.length - 1] as ChatModelResponse;
+	const newLastModelResponseItem = {...lastModelResponseItem};
+	newChatHistory[newChatHistory.length - 1] = newLastModelResponseItem;
+
+	const modelResponse = newLastModelResponseItem.response.slice();
+	newLastModelResponseItem.response = modelResponse;
+
+	const functionCall: ChatModelFunctionCall = {
+			type: "functionCall",
+			name: functionName,
+			description: functionDescription,
+			params: callParams,
+			result: callResult,
+			rawCall
+	};
+
+	if (startsNewChunk)
+			functionCall.startsNewChunk = true;
+
+	modelResponse.push(functionCall);
+
+	return newChatHistory;
+}
+
+export function createChatMessageArray(
+	messages: ChatMessage[],
+): ChatHistoryItem[] {
+	const items: ChatHistoryItem[] = []
+	let systemPrompt: string | undefined
+	for (const message of messages) {
+		if (message.role === 'user') {
+			items.push({
+				type: 'user',
+				text: flattenMessageTextContent(message.content),
+			})
+		} else if (message.role === 'assistant') {
+			items.push({
+				type: 'model',
+				response: [message.content],
+			})
+		} else if (message.role === 'system') {
+			if (systemPrompt) {
+				systemPrompt += '\n\n' + flattenMessageTextContent(message.content)
+			} else {
+				systemPrompt = flattenMessageTextContent(message.content)
+			}
+		}
+	}
+
+	if (systemPrompt) {
+		items.unshift({
+			type: 'system',
+			text: systemPrompt,
+		})
+	}
+
+	return items
+}
+
+export async function readGBNFFiles(directoryPath: string) {
+	const gbnfFiles = fs
+		.readdirSync(directoryPath)
+		.filter((f) => f.endsWith('.gbnf'))
+	const fileContents = await Promise.all(
+		gbnfFiles.map((file) =>
+			fs.promises.readFile(path.join(directoryPath, file), 'utf-8'),
+		),
+	)
+	return gbnfFiles.reduce((acc, file, i) => {
+		acc[file.replace('.gbnf', '')] = fileContents[i]
+		return acc
+	}, {} as Record<string, string>)
+}
+
+

package/src/engines/node-llama-cpp/validateModelFile.ts

@@ -0,0 +1,46 @@
+
+
+import fs from 'node:fs'
+import { calculateFileChecksum } from '#package/lib/calculateFileChecksum.js'
+import { resolveModelFileLocation } from '#package/lib/resolveModelFileLocation.js'
+
+interface ValidatableModelConfig {
+	url?: string
+	location?: string
+	modelsCachePath: string
+	sha256?: string
+}
+
+export async function validateModelFile(config: ValidatableModelConfig): Promise<string | undefined> {
+	const fileLocation = resolveModelFileLocation({
+		url: config.url,
+		filePath: config.location,
+		modelsCachePath: config.modelsCachePath,
+	})
+	if (!fs.existsSync(fileLocation)) {
+		return `Model 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 (config.sha256) {
+			const fileHash = await calculateFileChecksum(fileLocation, 'sha256')
+			if (fileHash === config.sha256) {
+				fs.unlinkSync(ipullFile)
+				validatedChecksum = true
+			}
+		}
+		if (!validatedChecksum) {
+			return `Model file with incomplete download`
+		}
+	}
+
+	if (!validatedChecksum && config.sha256) {
+		const fileHash = await calculateFileChecksum(fileLocation, 'sha256')
+		if (fileHash !== config.sha256) {
+			return `File sha256 checksum mismatch: expected ${config.sha256} got ${fileHash} at ${fileLocation}`
+		}
+	}
+	return undefined
+}

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

@@ -0,0 +1,369 @@
+import StableDiffusion from '@lmagder/node-stable-diffusion-cpp'
+import { gguf } from '@huggingface/gguf'
+import sharp from 'sharp'
+import fs from 'node:fs'
+import path from 'node:path'
+import {
+	EngineContext,
+	FileDownloadProgress,
+	ModelConfig,
+	EngineTextToImageResult,
+	ModelFileSource,
+	EngineTextToImageArgs,
+	Image,
+	EngineImageToImageArgs,
+} from '#package/types/index.js'
+import { LogLevel, LogLevels } from '#package/lib/logger.js'
+import { downloadModelFile } from '#package/lib/downloadModelFile.js'
+import { resolveModelFileLocation } from '#package/lib/resolveModelFileLocation.js'
+import { acquireFileLock } from '#package/lib/acquireFileLock.js'
+import { getRandomNumber } from '#package/lib/util.js'
+import { StableDiffusionSamplingMethod, StableDiffusionSchedule, StableDiffusionWeightType } from './types.js'
+import { validateModelFiles, ModelValidationResult } from './validateModelFiles.js'
+import { parseQuantization, getWeightType, getSamplingMethod } from './util.js'
+
+export interface StableDiffusionInstance {
+	context: StableDiffusion.Context
+}
+
+export interface StableDiffusionModelConfig extends ModelConfig {
+	location: string
+	sha256?: string
+	clipL?: ModelFileSource
+	clipG?: ModelFileSource
+	vae?: ModelFileSource
+	t5xxl?: ModelFileSource
+	controlNet?: ModelFileSource
+	taesd?: ModelFileSource
+	diffusionModel?: boolean
+	model?: ModelFileSource
+	loras?: ModelFileSource[]
+	samplingMethod?: StableDiffusionSamplingMethod
+	weightType?: StableDiffusionWeightType
+	schedule?: StableDiffusionSchedule
+	device?: {
+		gpu?: boolean | 'auto' | (string & {})
+		cpuThreads?: number
+	}
+}
+
+interface StableDiffusionModelMeta {
+	gguf: any
+}
+
+export const autoGpu = true
+
+export async function prepareModel(
+	{ config, log }: EngineContext<StableDiffusionModelConfig, StableDiffusionModelMeta>,
+	onProgress?: (progress: FileDownloadProgress) => void,
+	signal?: AbortSignal,
+) {
+	fs.mkdirSync(path.dirname(config.location), { recursive: true })
+	const releaseFileLock = await acquireFileLock(config.location)
+	if (signal?.aborted) {
+		releaseFileLock()
+		return
+	}
+	log(LogLevels.info, `Preparing stable-diffusion model at ${config.location}`, {
+		model: config.id,
+	})
+
+	const downloadModel = (url: string, validationResult: ModelValidationResult) => {
+		log(LogLevels.info, `${validationResult.message} - Downloading model files`, {
+			model: config.id,
+			url: config.url,
+			location: config.location,
+			errors: validationResult.errors,
+		})
+		const downloadPromises = []
+		if (validationResult.errors.model && config.location) {
+			downloadPromises.push(
+				downloadModelFile({
+					url: url,
+					filePath: config.location,
+					modelsCachePath: config.modelsCachePath,
+					onProgress,
+					signal,
+				}),
+			)
+		}
+		const pushDownload = (src: ModelFileSource) => {
+			if (!src.url) {
+				return
+			}
+			downloadPromises.push(
+				downloadModelFile({
+					url: src.url,
+					filePath: src.file,
+					modelsCachePath: config.modelsCachePath,
+					onProgress,
+					signal,
+				}),
+			)
+		}
+		if (validationResult.errors.clipG && config.clipG) {
+			pushDownload(config.clipG)
+		}
+		if (validationResult.errors.clipL && config.clipL) {
+			pushDownload(config.clipL)
+		}
+		if (validationResult.errors.vae && config.vae) {
+			pushDownload(config.vae)
+		}
+		if (validationResult.errors.t5xxl && config.t5xxl) {
+			pushDownload(config.t5xxl)
+		}
+		if (validationResult.errors.controlNet && config.controlNet) {
+			pushDownload(config.controlNet)
+		}
+		if (validationResult.errors.taesd && config.taesd) {
+			pushDownload(config.taesd)
+		}
+		if (config.loras) {
+			for (const lora of config.loras) {
+				if (!lora.url) {
+					continue
+				}
+				pushDownload(lora)
+			}
+		}
+		return Promise.all(downloadPromises)
+	}
+	try {
+		if (signal?.aborted) {
+			return
+		}
+
+		const validationResults = await validateModelFiles(config)
+		if (signal?.aborted) {
+			return
+		}
+		if (validationResults) {
+			if (config.url) {
+				await downloadModel(config.url, validationResults)
+			} else {
+				throw new Error(`${validationResults.message} - No URL provided`)
+			}
+		}
+
+		const finalValidationError = await validateModelFiles(config)
+		if (finalValidationError) {
+			throw new Error(`Downloaded files are invalid: ${finalValidationError}`)
+		}
+
+		const result: any = {}
+		if (config.location.endsWith('.gguf')) {
+			const { metadata, tensorInfos } = await gguf(config.location, {
+				allowLocalFile: true,
+			})
+			result.gguf = metadata
+		}
+		return result
+	} catch (error) {
+		throw error
+	} finally {
+		releaseFileLock()
+	}
+}
+
+export async function createInstance({ config, log }: EngineContext<StableDiffusionModelConfig>, signal?: AbortSignal) {
+	log(LogLevels.debug, 'Load Stable Diffusion model', config)
+	const handleLog = (level: string, message: string) => {
+		log(level as LogLevel, message)
+	}
+	const handleProgress = (step: number, steps: number, time: number) => {
+		log(LogLevels.debug, `Progress: ${step}/${steps} (${time}ms)`)
+	}
+
+	const resolveComponentLocation = (src?: ModelFileSource) => {
+		if (src) {
+			return resolveModelFileLocation({
+				url: src.url,
+				filePath: src.file,
+				modelsCachePath: config.modelsCachePath,
+			})
+		}
+		return undefined
+	}
+
+	const vaeFilePath = resolveComponentLocation(config.vae)
+	const clipLFilePath = resolveComponentLocation(config.clipL)
+	const clipGFilePath = resolveComponentLocation(config.clipG)
+	const t5xxlFilePath = resolveComponentLocation(config.t5xxl)
+	const controlNetFilePath = resolveComponentLocation(config.controlNet)
+	const taesdFilePath = resolveComponentLocation(config.taesd)
+
+	let weightType = config.weightType ? getWeightType(config.weightType) : undefined
+	if (typeof weightType === 'undefined') {
+		const quantization = parseQuantization(config.location)
+		if (quantization) {
+			weightType = getWeightType(quantization)
+		}
+	}
+
+	if (typeof weightType === 'undefined') {
+		log(LogLevels.warn, 'Failed to parse model weight type (quantization) from file name, falling back to f32', {
+			file: config.location,
+		})
+	}
+
+	const loraDir = path.join(path.dirname(config.location), 'loras')
+	const contextParams = {
+		model: !config.diffusionModel ? config.location : undefined,
+		diffusionModel: config.diffusionModel ? config.location : undefined,
+		numThreads: config.device?.cpuThreads,
+		vae: vaeFilePath,
+		clipL: clipLFilePath,
+		clipG: clipGFilePath,
+		t5xxl: t5xxlFilePath,
+		controlNet: controlNetFilePath,
+		taesd: taesdFilePath,
+		weightType: weightType,
+		loraDir: loraDir,
+		// TODO how to expose?
+		// keepClipOnCpu: true,
+		// keepControlNetOnCpu: true,
+		// keepVaeOnCpu: true,
+	}
+	log(LogLevels.debug, 'Creating context with', contextParams)
+	const context = await StableDiffusion.createContext(
+		// @ts-ignore
+		contextParams,
+		handleLog,
+		handleProgress,
+	)
+
+	return {
+		context,
+	}
+}
+
+export async function processTextToImageTask(
+	{ request, config, log }: EngineTextToImageArgs<StableDiffusionModelConfig>,
+	instance: StableDiffusionInstance,
+	signal?: AbortSignal,
+): Promise<EngineTextToImageResult> {
+	const seed = request.seed ?? getRandomNumber(0, 1000000)
+	const results = await instance.context.txt2img({
+		prompt: request.prompt,
+		negativePrompt: request.negativePrompt,
+		width: request.width || 512,
+		height: request.height || 512,
+		batchCount: request.batchCount,
+		sampleMethod: getSamplingMethod(request.samplingMethod || config.samplingMethod),
+		sampleSteps: request.sampleSteps,
+		cfgScale: request.cfgScale,
+		// @ts-ignore
+		guidance: request.guidance,
+		styleRatio: request.styleRatio,
+		controlStrength: request.controlStrength,
+		normalizeInput: false,
+		seed,
+	})
+
+	const images: Image[] = []
+	for (const [idx, img] of results.entries()) {
+		images.push({
+			handle: sharp(img.data, {
+				raw: {
+					width: img.width,
+					height: img.height,
+					channels: img.channel,
+				},
+			}),
+			width: img.width,
+			height: img.height,
+			channels: img.channel,
+		})
+	}
+	if (!images.length) {
+		throw new Error('No images generated')
+	}
+	return {
+		images: images,
+		seed,
+	}
+}
+
+export async function processImageToImageTask(
+	{ request, config, log }: EngineImageToImageArgs<StableDiffusionModelConfig>,
+	instance: StableDiffusionInstance,
+	signal?: AbortSignal,
+): Promise<EngineTextToImageResult> {
+	const seed = request.seed ?? getRandomNumber(0, 1000000)
+	console.debug('processImageToImageTask', {
+		width: request.image.width,
+		height: request.image.height,
+		channel: request.image.channels as 3 | 4,
+	})
+	const initImage = {
+		data: await request.image.handle.raw().toBuffer(),
+		width: request.image.width,
+		height: request.image.height,
+		channel: request.image.channels as 3 | 4,
+	}
+	const results = await instance.context.img2img({
+		initImage,
+		prompt: request.prompt,
+		width: request.width || 512,
+		height: request.height || 512,
+		batchCount: request.batchCount,
+		sampleMethod: getSamplingMethod(request.samplingMethod || config.samplingMethod),
+		cfgScale: request.cfgScale,
+		sampleSteps: request.sampleSteps,
+		// @ts-ignore
+		guidance: request.guidance,
+		strength: request.strength,
+		styleRatio: request.styleRatio,
+		controlStrength: request.controlStrength,
+		seed,
+	})
+
+	const images: Image[] = []
+	// to sharp
+	// const imagePromises = results.map(async (img, idx) => {
+	// 	return await sharp(img.data, {
+	// 			raw: {
+	// 				width: img.width,
+	// 				height: img.height,
+	// 				channels: img.channel,
+	// 			},
+	// 		})
+	// 	})
+
+	for (const [idx, img] of results.entries()) {
+		console.debug('img', {
+			id: idx,
+			width: img.width,
+			height: img.height,
+			channels: img.channel,
+		})
+
+		images.push({
+			handle: sharp(img.data, {
+				raw: {
+					width: img.width,
+					height: img.height,
+					channels: img.channel,
+				},
+			}),
+			width: img.width,
+			height: img.height,
+			channels: img.channel,
+		})
+
+		// images.push({
+		// 	data: img.data,
+		// 	width: img.width,
+		// 	height: img.height,
+		// 	channels: img.channel,
+		// })
+	}
+	if (!images.length) {
+		throw new Error('No images generated')
+	}
+	return {
+		images: images,
+		seed,
+	}
+}

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

@@ -0,0 +1,54 @@
+export type StableDiffusionWeightType =
+	| 'f32'
+	| 'f16'
+	| 'q4_0'
+	| 'q4_1'
+	| 'q5_0'
+	| 'q5_1'
+	| 'q8_0'
+	| 'q2_k'
+	| 'q3_k'
+	| 'q4_k'
+	| 'q5_k'
+	| 'q6_k'
+	| 'q8_k'
+	| 'iq2_xxs'
+	| 'iq2_xs'
+	| 'iq3_xxs'
+	| 'iq1_s'
+	| 'iq4_nl'
+	| 'iq3_s'
+	| 'iq2_s'
+	| 'iq4_xs'
+	| 'i8'
+	| 'i16'
+	| 'i32'
+	| 'i64'
+	| 'f64'
+	| 'iq1_m'
+	| 'bf16'
+	| 'q4_0_4_4'
+	| 'q4_0_4_8'
+	| 'q4_0_8_8'
+	| (string & {})
+
+export type StableDiffusionSchedule =
+	| 'discrete'
+	| 'karras'
+	| 'exponential'
+	| 'ays'
+	| 'gits'
+	| (string & {})
+
+export type StableDiffusionSamplingMethod =
+	| 'euler'
+	| 'euler_a'
+	| 'lcm'
+	| 'heun'
+	| 'dpm2'
+	| 'dpm++2s_a'
+	| 'dpm++2m'
+	| 'dpm++2mv2'
+	| 'ipndm'
+	| 'ipndm_v'
+	| (string & {})

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

@@ -0,0 +1,58 @@
+import StableDiffusion from '@lmagder/node-stable-diffusion-cpp'
+
+export function parseQuantization(filename: string): string | null {
+	// Regular expressions to match different quantization patterns
+	const regexPatterns = [
+		/q(\d+)_(\d+)/i, // q4_0
+		/[-_\.](f16|f32|int8|int4)/i, // f16, f32, int8, int4
+		/[-_\.](fp16|fp32)/i, // fp16, fp32
+	]
+
+	for (const regex of regexPatterns) {
+		const match = filename.match(regex)
+		if (match) {
+			// If there's a match, return the full matched quantization string
+			// Remove leading dash if present, convert to uppercase
+			return match[0].replace(/^[-_]/, '').replace(/fp/i, 'f').toLowerCase()
+		}
+	}
+	return null
+}
+
+export function getWeightType(key: string): number | undefined {
+	const weightKey = key.toUpperCase() as keyof typeof StableDiffusion.Type
+	if (weightKey in StableDiffusion.Type) {
+		return StableDiffusion.Type[weightKey]
+	}
+	console.warn('Unknown weight type', weightKey)
+	return undefined
+}
+
+export function getSamplingMethod(method?: string): StableDiffusion.SampleMethod | undefined {
+	switch (method) {
+		case 'euler':
+			return StableDiffusion.SampleMethod.Euler
+		case 'euler_a':
+			return StableDiffusion.SampleMethod.EulerA
+		case 'lcm':
+			return StableDiffusion.SampleMethod.LCM
+		case 'heun':
+			return StableDiffusion.SampleMethod.Heun
+		case 'dpm2':
+			return StableDiffusion.SampleMethod.DPM2
+		case 'dpm++2s_a':
+			return StableDiffusion.SampleMethod.DPMPP2SA
+		case 'dpm++2m':
+			return StableDiffusion.SampleMethod.DPMPP2M
+		case 'dpm++2mv2':
+			return StableDiffusion.SampleMethod.DPMPP2Mv2
+		case 'ipndm':
+			// @ts-ignore
+			return StableDiffusion.SampleMethod.IPNDM
+		case 'ipndm_v':
+			// @ts-ignore
+			return StableDiffusion.SampleMethod.IPNDMV
+	}
+	console.warn('Unknown sampling method', method)
+	return undefined
+}