From f6d15148e9e7b494b04931f1c948ef5a218b5d91 Mon Sep 17 00:00:00 2001 From: james-333i Date: Tue, 25 Aug 2026 14:36:42 -0700 Subject: [PATCH] Add vision support to LlamaLanguageModel via mtmd Image segments threw unsupportedFeature because the backend had no multimodal path, even though the prebuilt llama.cpp binaries ship the mtmd library and its helpers. Accept an mmprojPath at initialization and load the projector next to the model. When a projector is present, prompt formatting replaces each image segment with the mtmd media marker and collects payloads in order, then generation tokenizes the marker-annotated prompt with mtmd_tokenize and evaluates text and image chunks through mtmd_helper_eval_chunks before sampling continues from the resulting position. Both respond and streaming support images, and models without a projector keep rejecting image input. Adds live tests generating from an embedded test image through both paths. --- .../Models/LlamaLanguageModel.swift | 353 ++++++++++++++++-- .../LlamaLanguageModelTests.swift | 75 ++++ 2 files changed, 397 insertions(+), 31 deletions(-) diff --git a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift index d2f29b03..4dd41aa0 100644 --- a/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift +++ b/Sources/AnyLanguageModel/Models/LlamaLanguageModel.swift @@ -208,6 +208,11 @@ import Foundation /// A negative value offloads all layers, and `0` runs entirely on the CPU. public let gpuLayers: Int32 + /// The path to the multimodal projector GGUF file, when the model has one. + /// + /// Prompts may include image segments only when a projector is loaded. + public let mmprojPath: String? + /// The default GPU layer count for the current platform. /// /// All layers are offloaded by default: the prebuilt llama.cpp binaries @@ -453,6 +458,9 @@ import Foundation /// The model's vocabulary private var vocab: OpaquePointer? + /// The multimodal projector context, when a projector file was provided + private var mtmdContext: OpaquePointer? + /// Whether the model is currently loaded private var isModelLoaded: Bool = false @@ -462,9 +470,18 @@ import Foundation /// - modelPath: The path to the GGUF model file. /// - gpuLayers: The number of model layers to offload to the GPU. /// Defaults to ``defaultGPULayerCount``. - public init(modelPath: String, gpuLayers: Int32 = LlamaLanguageModel.defaultGPULayerCount) { + /// - mmprojPath: The path to a multimodal projector GGUF file matching + /// the model. When provided, prompts may include image segments, + /// which are encoded through the projector. Defaults to `nil` + /// (text only). + public init( + modelPath: String, + gpuLayers: Int32 = LlamaLanguageModel.defaultGPULayerCount, + mmprojPath: String? = nil + ) { self.modelPath = modelPath self.gpuLayers = gpuLayers + self.mmprojPath = mmprojPath self.legacyDefaults = ResolvedGenerationOptions() } @@ -505,6 +522,9 @@ import Foundation } deinit { + if let mtmdContext = mtmdContext { + mtmd_free(mtmdContext) + } if let model = model { llama_model_free(model) } @@ -517,8 +537,9 @@ import Foundation includeSchemaInPrompt: Bool, options: GenerationOptions ) async throws -> LanguageModelSession.Response where Content: Generable { - // Validate that no image segments are present - try validateNoImageSegments(in: session) + if mmprojPath == nil { + try validateNoImageSegments(in: session) + } try await ensureModelLoaded() let runtimeOptions = resolvedOptions(from: options) @@ -542,6 +563,9 @@ import Foundation llama_set_causal_attn(context, true) llama_set_n_threads(context, runtimeOptions.threads, runtimeOptions.threads) + var promptImages: [Data] = [] + let imageMarker = mtmdContext != nil ? String(cString: mtmd_default_marker()) : nil + let fullPrompt: String if includeSchemaInPrompt, type != String.self { fullPrompt = try formatPrompt( @@ -550,18 +574,39 @@ import Foundation assistantPrefill: runtimeOptions.assistantPrefill ) } else { - fullPrompt = try formatPrompt(for: session, assistantPrefill: runtimeOptions.assistantPrefill) + fullPrompt = try formatPrompt( + for: session, + extraSystemMessage: nil, + assistantPrefill: runtimeOptions.assistantPrefill, + imageMarker: imageMarker, + images: &promptImages + ) } if type == String.self { let maxTokens = runtimeOptions.maximumResponseTokens ?? 100 - let text = try await generateText( - context: context, - model: model!, - prompt: fullPrompt, - maxTokens: maxTokens, - options: runtimeOptions - ) + let text: String + if promptImages.isEmpty { + text = try await generateText( + context: context, + model: model!, + prompt: fullPrompt, + maxTokens: maxTokens, + options: runtimeOptions + ) + } else { + var accumulated = "" + try performMultimodalGeneration( + context: context, + prompt: fullPrompt, + images: promptImages, + maxTokens: maxTokens, + options: runtimeOptions + ) { tokenText in + accumulated += tokenText + } + text = accumulated + } return LanguageModelSession.Response( content: text as! Content, @@ -599,15 +644,16 @@ import Foundation fatalError("LlamaLanguageModel only supports generating String content") } - // Validate that no image segments are present - do { - try validateNoImageSegments(in: session) - } catch { - return LanguageModelSession.ResponseStream( - stream: AsyncThrowingStream { continuation in - continuation.finish(throwing: error) - } - ) + if mmprojPath == nil { + do { + try validateNoImageSegments(in: session) + } catch { + return LanguageModelSession.ResponseStream( + stream: AsyncThrowingStream { continuation in + continuation.finish(throwing: error) + } + ) + } } let stream: AsyncThrowingStream.Snapshot, any Error> = @@ -637,18 +683,18 @@ import Foundation llama_set_n_threads(context, runtimeOptions.threads, runtimeOptions.threads) var accumulatedText = "" + var promptImages: [Data] = [] + let imageMarker = + self.mtmdContext != nil ? String(cString: mtmd_default_marker()) : nil let fullPrompt = try self.formatPrompt( for: session, - assistantPrefill: runtimeOptions.assistantPrefill + extraSystemMessage: nil, + assistantPrefill: runtimeOptions.assistantPrefill, + imageMarker: imageMarker, + images: &promptImages ) - try self.performTextGeneration( - context: context, - model: model!, - prompt: fullPrompt, - maxTokens: maxTokens, - options: runtimeOptions - ) { tokenText in + let yieldToken: (String) -> Void = { tokenText in accumulatedText += tokenText let snapshot = LanguageModelSession.ResponseStream.Snapshot( @@ -658,6 +704,26 @@ import Foundation continuation.yield(snapshot) } + if promptImages.isEmpty { + try self.performTextGeneration( + context: context, + model: model!, + prompt: fullPrompt, + maxTokens: maxTokens, + options: runtimeOptions, + onToken: yieldToken + ) + } else { + try self.performMultimodalGeneration( + context: context, + prompt: fullPrompt, + images: promptImages, + maxTokens: maxTokens, + options: runtimeOptions, + onToken: yieldToken + ) + } + continuation.finish() } catch { continuation.finish(throwing: error) @@ -686,6 +752,10 @@ import Foundation llama_backend_init() // Free any existing model before loading a new one + if let existingContext = mtmdContext { + mtmd_free(existingContext) + self.mtmdContext = nil + } if let existingModel = model { llama_model_free(existingModel) self.model = nil @@ -696,6 +766,22 @@ import Foundation throw LlamaLanguageModelError.modelLoadFailed } + if let mmprojPath { + guard FileManager.default.fileExists(atPath: mmprojPath) else { + llama_model_free(loadedModel) + throw LlamaLanguageModelError.invalidModelPath + } + var mtmdParams = mtmd_context_params_default() + mtmdParams.use_gpu = gpuLayers != 0 + mtmdParams.print_timings = false + mtmdParams.n_threads = legacyDefaults.threads + guard let projector = mtmd_init_from_file(mmprojPath, loadedModel, mtmdParams) else { + llama_model_free(loadedModel) + throw LlamaLanguageModelError.modelLoadFailed + } + self.mtmdContext = projector + } + self.model = loadedModel self.vocab = llama_model_get_vocab(loadedModel) self.isModelLoaded = true @@ -1285,6 +1371,146 @@ import Foundation } } + /// Evaluates a marker-annotated multimodal prompt through the projector, + /// then generates text tokens from the resulting state. + private func performMultimodalGeneration( + context: OpaquePointer, + prompt: String, + images: [Data], + maxTokens: Int, + options: ResolvedGenerationOptions, + onToken: (String) -> Void + ) throws { + guard let mtmdContext, let model = self.model, + let vocab = llama_model_get_vocab(model) + else { + throw LlamaLanguageModelError.contextInitializationFailed + } + + var bitmaps: [OpaquePointer?] = [] + defer { + for bitmap in bitmaps { + if let bitmap { + mtmd_bitmap_free(bitmap) + } + } + } + for imageData in images { + // Pinned to the current llama.swift signature. llama.cpp master adds a + // trailing options argument to this helper; update alongside the dependency. + let wrapper = imageData.withUnsafeBytes { raw -> mtmd_helper_bitmap_wrapper in + mtmd_helper_bitmap_init_from_buf( + mtmdContext, + raw.bindMemory(to: UInt8.self).baseAddress, + imageData.count, + false + ) + } + if let videoContext = wrapper.video_ctx { + mtmd_helper_video_free(videoContext) + throw LlamaLanguageModelError.unsupportedFeature + } + guard let bitmap = wrapper.bitmap else { + throw LlamaLanguageModelError.encodingFailed + } + bitmaps.append(bitmap) + } + + guard let chunks = mtmd_input_chunks_init() else { + throw LlamaLanguageModelError.encodingFailed + } + defer { mtmd_input_chunks_free(chunks) } + + let tokenizeResult = prompt.withCString { cPrompt -> Int32 in + var inputText = mtmd_input_text( + text: cPrompt, + text_len: strlen(cPrompt), + add_special: true, + parse_special: true + ) + return bitmaps.withUnsafeMutableBufferPointer { buffer in + mtmd_tokenize(mtmdContext, chunks, &inputText, buffer.baseAddress, buffer.count) + } + } + guard tokenizeResult == 0 else { + throw LlamaLanguageModelError.tokenizationFailed + } + + var pastPosition: llama_pos = 0 + let evalResult = mtmd_helper_eval_chunks( + mtmdContext, + context, + chunks, + 0, + 0, + Int32(options.batchSize), + true, + &pastPosition + ) + guard evalResult == 0 else { + throw LlamaLanguageModelError.decodingFailed + } + + guard let sampler = llama_sampler_chain_init(llama_sampler_chain_default_params()) else { + throw LlamaLanguageModelError.decodingFailed + } + defer { llama_sampler_free(sampler) } + let samplerPtr = UnsafeMutablePointer(sampler) + + if options.repeatPenalty != 1.0 || options.frequencyPenalty != 0.0 || options.presencePenalty != 0.0 { + llama_sampler_chain_add( + samplerPtr, + llama_sampler_init_penalties( + llama_vocab_n_tokens(vocab), + options.repeatLastN, + options.repeatPenalty, + options.frequencyPenalty, + options.presencePenalty + ) + ) + } + applySampling(sampler: samplerPtr, effectiveTemperature: options.temperature, options: options) + + var batch = llama_batch_init(1, 0, 1) + defer { llama_batch_free(batch) } + + var n_cur: Int32 = Int32(pastPosition) + var sampleIndex: Int32 = -1 + + for _ in 0 ..< maxTokens { + if Task.isCancelled { + break + } + + let nextToken = llama_sampler_sample(samplerPtr, context, sampleIndex) + llama_sampler_accept(samplerPtr, nextToken) + + if llama_vocab_is_eog(vocab, nextToken) { + break + } + + if let tokenText = tokenToText(vocab: vocab, token: nextToken) { + onToken(tokenText) + } + + batch.n_tokens = 1 + batch.token[0] = nextToken + batch.pos[0] = n_cur + batch.n_seq_id[0] = 1 + if let seq_ids = batch.seq_id, let seq_id = seq_ids[0] { + seq_id[0] = 0 + } + batch.logits[0] = 1 + + n_cur += 1 + + guard llama_decode(context, batch) == 0 else { + break + } + sampleIndex = 0 + } + } + // MARK: - Image Validation private func validateNoImageSegments(in session: LanguageModelSession) throws { @@ -1422,6 +1648,23 @@ import Foundation for session: LanguageModelSession, extraSystemMessage: String? = nil, assistantPrefill: String? = nil + ) throws -> String { + var images: [Data] = [] + return try formatPrompt( + for: session, + extraSystemMessage: extraSystemMessage, + assistantPrefill: assistantPrefill, + imageMarker: nil, + images: &images + ) + } + + private func formatPrompt( + for session: LanguageModelSession, + extraSystemMessage: String?, + assistantPrefill: String?, + imageMarker: String?, + images: inout [Data] ) throws -> String { guard let model = self.model else { throw LlamaLanguageModelError.modelLoadFailed @@ -1432,19 +1675,31 @@ import Foundation for entry in session.transcript { switch entry { case .instructions(let instructions): - let text = extractText(from: instructions.segments) + let text = try extractContent( + from: instructions.segments, + imageMarker: imageMarker, + images: &images + ) if !text.isEmpty { messages.append(("system", text)) } case .prompt(let prompt): - let text = extractText(from: prompt.segments) + let text = try extractContent( + from: prompt.segments, + imageMarker: imageMarker, + images: &images + ) if !text.isEmpty { messages.append(("user", text)) } case .response(let response): - let text = extractText(from: response.segments) + let text = try extractContent( + from: response.segments, + imageMarker: imageMarker, + images: &images + ) if !text.isEmpty { messages.append(("assistant", text)) } @@ -1522,6 +1777,42 @@ import Foundation }.joined() } + /// Extracts message content from segments, replacing each image segment + /// with `imageMarker` and collecting its payload in order. Image segments + /// throw ``LlamaLanguageModelError/unsupportedFeature`` when no marker is + /// provided. + private func extractContent( + from segments: [Transcript.Segment], + imageMarker: String?, + images: inout [Data] + ) throws -> String { + var parts: [String] = [] + for segment in segments { + switch segment { + case .text(let t): + parts.append(t.content) + case .image(let image): + guard let imageMarker else { + throw LlamaLanguageModelError.unsupportedFeature + } + switch image.source { + case .data(let data, _): + images.append(data) + parts.append(imageMarker) + case .url(let url): + guard url.isFileURL, let data = try? Data(contentsOf: url) else { + throw LlamaLanguageModelError.unsupportedFeature + } + images.append(data) + parts.append(imageMarker) + } + default: + break + } + } + return parts.joined() + } + private func tokenizeText(vocab: OpaquePointer, text: String) throws -> [llama_token] { let utf8Count = text.utf8.count let maxTokens = Int32(max(utf8Count * 2, 8)) // Rough estimate, minimum capacity diff --git a/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift b/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift index a319b82c..ed5a18f6 100644 --- a/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift +++ b/Tests/AnyLanguageModelTests/LlamaLanguageModelTests.swift @@ -425,3 +425,78 @@ import Testing } } #endif // Llama + +#if Llama + @Suite( + "LlamaLanguageModel vision", + .serialized, + .enabled( + if: ProcessInfo.processInfo.environment["LLAMA_VISION_MODEL_PATH"] != nil + && ProcessInfo.processInfo.environment["LLAMA_VISION_MMPROJ_PATH"] != nil + ) + ) + struct LlamaLanguageModelVisionTests { + static let redSquarePNG = Data( + base64Encoded: "iVBORw0KGgoAAAANSUhEUgAAAGAAAABgCAIAAABt+uBvAAABC0lEQVR4nO3OMQ0AIAAEsfdvGhyw9gaS" + + "CujO9j34QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E" + + "+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E" + + "+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E" + + "+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E+UGcH8T5QZwfxPlBnB/E" + + "+UHcBWwZ3g5gacwjAAAAAElFTkSuQmCC" + )! + + let model = LlamaLanguageModel( + modelPath: ProcessInfo.processInfo.environment["LLAMA_VISION_MODEL_PATH"]!, + mmprojPath: ProcessInfo.processInfo.environment["LLAMA_VISION_MMPROJ_PATH"]! + ) + + @Test func describesImageData() async throws { + let transcript = Transcript(entries: [ + .prompt( + Transcript.Prompt(segments: [ + .text(.init(content: "What is the dominant color of this image? Answer with one word.")), + .image(.init(data: Self.redSquarePNG, mimeType: "image/png")), + ]) + ) + ]) + let session = LanguageModelSession(model: model, transcript: transcript) + let response = try await session.respond(to: "") + #expect(response.content.lowercased().contains("red")) + } + + @Test func streamsImageDescription() async throws { + let transcript = Transcript(entries: [ + .prompt( + Transcript.Prompt(segments: [ + .text(.init(content: "What is the dominant color of this image? Answer with one word.")), + .image(.init(data: Self.redSquarePNG, mimeType: "image/png")), + ]) + ) + ]) + let session = LanguageModelSession(model: model, transcript: transcript) + let stream = session.streamResponse(to: "") + var last = "" + for try await snapshot in stream { + last = snapshot.content + } + #expect(last.lowercased().contains("red")) + } + + @Test func rejectsImagesWithoutProjector() async throws { + let textOnlyModel = LlamaLanguageModel( + modelPath: ProcessInfo.processInfo.environment["LLAMA_VISION_MODEL_PATH"]! + ) + let transcript = Transcript(entries: [ + .prompt( + Transcript.Prompt(segments: [ + .image(.init(data: Self.redSquarePNG, mimeType: "image/png")) + ]) + ) + ]) + let session = LanguageModelSession(model: textOnlyModel, transcript: transcript) + await #expect(throws: LlamaLanguageModelError.self) { + _ = try await session.respond(to: "") + } + } + } +#endif