Working on a solution

Author: DePasqualeOrg

Applications/VLMEval/ContentView.swift

@@ -331,7 +331,16 @@ class VLMEvaluator {
             MLXRandom.seed(UInt64(Date.timeIntervalSinceReferenceDate * 1000))
 
             let result = try await modelContainer.perform { context in
-                var userInput = UserInput(prompt: prompt, images: [.ciImage(image)])
+                var userInput = UserInput(
+                    messages: [
+                        [
+                            "role": "user",
+                            "content": [
+                                ["type": "text", "text": prompt],
+                                ["type": "image"],
+                            ],
+                        ]
+                    ], images: [.ciImage(image)])
                 userInput.processing.resize = .init(width: 448, height: 448)
 
                 let input = try await context.processor.prepare(input: userInput)

Libraries/MLXLMCommon/UserInput.swift

@@ -4,18 +4,19 @@ import CoreImage
 import Foundation
 import MLX
 
+public typealias Message = [String: Any]
+
 /// Container for raw user input.
 ///
 /// A ``UserInputProcessor`` can convert this to ``LMInput``.
 /// See also ``ModelContext``.
 public struct UserInput: Sendable {
-
     /// Representation of a prompt or series of messages (conversation).
     public enum Prompt: Sendable, CustomStringConvertible {
         case text(String)
-        case messages([[String: String]])
+        case messages([Message])
 
-        public func asMessages() -> [[String: String]] {
+        public func asMessages() -> [Message] {
             switch self {
             case .text(let text):
                 return [["role": "user", "content": text]]
@@ -116,7 +117,7 @@ public struct UserInput: Sendable {
         self.images = images
     }
 
-    public init(messages: [[String: String]], images: [Image] = [Image]()) {
+    public init(messages: [Message], images: [Image] = [Image]()) {
         self.prompt = .messages(messages)
         self.images = images
     }

Libraries/MLXVLM/Models/Idefics3.swift

@@ -805,7 +805,7 @@ public class Idefics3Processor: UserInputProcessor {
     }
 
     public func prepare(input: UserInput) throws -> LMInput {
-        let prompt = input.prompt.asMessages().last?["content"] ?? ""
+        let prompt = input.prompt.asMessages().last?["content"] as? String ?? ""
 
         if input.images.isEmpty {
             // No image scenario

Libraries/MLXVLM/Models/Paligemma.swift

@@ -478,7 +478,7 @@ public class PaligGemmaProcessor: UserInputProcessor {
         }
 
         // this doesn't have a chat template so just use the last message.
-        var prompt = input.prompt.asMessages().last?["content"] ?? ""
+        var prompt = input.prompt.asMessages().last?["content"] as? String ?? ""
 
         // based on transformers/processing_paligemma
         let count = input.images.count * config.imageSequenceLength

Libraries/MLXVLM/Models/Qwen2VL.swift

@@ -686,24 +686,25 @@ public class Qwen2VLProcessor: UserInputProcessor {
         return (flattenedPatches, .init(gridT, gridH, gridW))
     }
 
-    private func prepareMessages(_ messages: [[String: String]], imageTHW: [THW]?) -> [[String: String]] {
+    private func prepareMessages(_ messages: [Message], imageTHW: [THW]?) -> [Message] {
         var messages = messages
+        print(messages)
         // Add system message if not present
-        if messages[0]["role"] != "system" {
+        if let role = messages[0]["role"] as? String, role != "system" {
             messages.insert(["role": "system", "content": "You are a helpful assistant."], at: 0)
         }
-        // Add image markers to last message if needed
-        if let imageTHW {
-            let lastIndex = messages.count - 1
-            var content = messages[lastIndex]["content"] ?? ""
-            let mergeLength = config.mergeSize * config.mergeSize
-            for thw in imageTHW {
-                content += "<|vision_start|>"
-                content += Array(repeating: "<|image_pad|>", count: thw.product / mergeLength).joined()
-                content += "<|vision_end|>"
-            }
-            messages[lastIndex]["content"] = content
-        }
+        //        // Add image markers to last message if needed
+        //        if let imageTHW {
+        //            let lastIndex = messages.count - 1
+        //            var content = messages[lastIndex]["content"] ?? ""
+        //            let mergeLength = config.mergeSize * config.mergeSize
+        //            for thw in imageTHW {
+        //                content += "<|vision_start|>"
+        //                content += Array(repeating: "<|image_pad|>", count: thw.product / mergeLength).joined()
+        //                content += "<|vision_end|>"
+        //            }
+        //            messages[lastIndex]["content"] = content
+        //        }
         return messages
     }