Merge branch 'main' into qwen-2.5-vl

DePasqualeOrg · web-flow · commit 362939f45bd1 · 2025-03-10T13:55:13.000+01:00
diff --git a/Libraries/Embedders/Models.swift b/Libraries/Embedders/Models.swift
@@ -83,6 +83,12 @@ public struct ModelConfiguration: Sendable {
             return ModelConfiguration(id: id)
         }
     }
+
+    @MainActor
+    public static var models: some Collection<ModelConfiguration> & Sendable {
+        bootstrap()
+        return Self.registry.values
+    }
 }
 
 extension ModelConfiguration {
diff --git a/Libraries/MLXLLM/LLMModelFactory.swift b/Libraries/MLXLLM/LLMModelFactory.swift
@@ -22,25 +22,42 @@ private func create<C: Codable, M>(
 /// Typically called via ``LLMModelFactory/load(hub:configuration:progressHandler:)``.
 public class ModelTypeRegistry: @unchecked Sendable {
 
+    /// Creates an empty registry.
+    public init() {
+        self.creators = [:]
+    }
+
+    /// Creates a registry with given creators.
+    public init(creators: [String: @Sendable (URL) throws -> any LanguageModel]) {
+        self.creators = creators
+    }
+
+    /// Shared instance with default model types.
+    public static let shared: ModelTypeRegistry = .init(creators: all())
+
+    /// All predefined model types.
+    private static func all() -> [String: @Sendable (URL) throws -> any LanguageModel] {
+        [
+            "mistral": create(LlamaConfiguration.self, LlamaModel.init),
+            "llama": create(LlamaConfiguration.self, LlamaModel.init),
+            "phi": create(PhiConfiguration.self, PhiModel.init),
+            "phi3": create(Phi3Configuration.self, Phi3Model.init),
+            "phimoe": create(PhiMoEConfiguration.self, PhiMoEModel.init),
+            "gemma": create(GemmaConfiguration.self, GemmaModel.init),
+            "gemma2": create(Gemma2Configuration.self, Gemma2Model.init),
+            "qwen2": create(Qwen2Configuration.self, Qwen2Model.init),
+            "starcoder2": create(Starcoder2Configuration.self, Starcoder2Model.init),
+            "cohere": create(CohereConfiguration.self, CohereModel.init),
+            "openelm": create(OpenElmConfiguration.self, OpenELMModel.init),
+            "internlm2": create(InternLM2Configuration.self, InternLM2Model.init),
+        ]
+    }
+
     // Note: using NSLock as we have very small (just dictionary get/set)
     // critical sections and expect no contention.  this allows the methods
     // to remain synchronous.
     private let lock = NSLock()
-
-    private var creators: [String: @Sendable (URL) throws -> any LanguageModel] = [
-        "mistral": create(LlamaConfiguration.self, LlamaModel.init),
-        "llama": create(LlamaConfiguration.self, LlamaModel.init),
-        "phi": create(PhiConfiguration.self, PhiModel.init),
-        "phi3": create(Phi3Configuration.self, Phi3Model.init),
-        "phimoe": create(PhiMoEConfiguration.self, PhiMoEModel.init),
-        "gemma": create(GemmaConfiguration.self, GemmaModel.init),
-        "gemma2": create(Gemma2Configuration.self, Gemma2Model.init),
-        "qwen2": create(Qwen2Configuration.self, Qwen2Model.init),
-        "starcoder2": create(Starcoder2Configuration.self, Starcoder2Model.init),
-        "cohere": create(CohereConfiguration.self, CohereModel.init),
-        "openelm": create(OpenElmConfiguration.self, OpenELMModel.init),
-        "internlm2": create(InternLM2Configuration.self, InternLM2Model.init),
-    ]
+    private var creators: [String: @Sendable (URL) throws -> any LanguageModel]
 
     /// Add a new model to the type registry.
     public func registerModelType(
@@ -72,8 +89,21 @@ public class ModelTypeRegistry: @unchecked Sendable {
 /// implementation, if needed.
 public class ModelRegistry: @unchecked Sendable {
 
+    /// Creates an empty registry.
+    public init() {
+        self.registry = Dictionary()
+    }
+
+    /// Creates a new registry with from given model configurations.
+    public init(modelConfigurations: [ModelConfiguration]) {
+        self.registry = Dictionary(uniqueKeysWithValues: modelConfigurations.map { ($0.name, $0) })
+    }
+
+    /// Shared instance with default model configurations.
+    public static let shared = ModelRegistry(modelConfigurations: all())
+
     private let lock = NSLock()
-    private var registry = Dictionary(uniqueKeysWithValues: all().map { ($0.name, $0) })
+    private var registry: [String: ModelConfiguration]
 
     static public let smolLM_135M_4bit = ModelConfiguration(
         id: "mlx-community/SmolLM-135M-Instruct-4bit",
@@ -226,6 +256,12 @@ public class ModelRegistry: @unchecked Sendable {
             }
         }
     }
+
+    public var models: some Collection<ModelConfiguration> & Sendable {
+        lock.withLock {
+            return registry.values
+        }
+    }
 }
 
 private struct LLMUserInputProcessor: UserInputProcessor {
@@ -268,13 +304,19 @@ private struct LLMUserInputProcessor: UserInputProcessor {
 /// ```
 public class LLMModelFactory: ModelFactory {
 
-    public static let shared = LLMModelFactory()
+    public init(typeRegistry: ModelTypeRegistry, modelRegistry: ModelRegistry) {
+        self.typeRegistry = typeRegistry
+        self.modelRegistry = modelRegistry
+    }
+
+    /// Shared instance with default behavior.
+    public static let shared = LLMModelFactory(typeRegistry: .shared, modelRegistry: .shared)
 
     /// registry of model type, e.g. configuration value `llama` -> configuration and init methods
-    public let typeRegistry = ModelTypeRegistry()
+    public let typeRegistry: ModelTypeRegistry
 
     /// registry of model id to configuration, e.g. `mlx-community/Llama-3.2-3B-Instruct-4bit`
-    public let modelRegistry = ModelRegistry()
+    public let modelRegistry: ModelRegistry
 
     public func configuration(id: String) -> ModelConfiguration {
         modelRegistry.configuration(id: id)
diff --git a/Libraries/MLXVLM/VLMModelFactory.swift b/Libraries/MLXVLM/VLMModelFactory.swift
@@ -67,17 +67,34 @@ private func create<C: Codable, P>(
 /// Typically called via ``LLMModelFactory/load(hub:configuration:progressHandler:)``.
 public class ModelTypeRegistry: @unchecked Sendable {
 
+    /// Creates an empty registry.
+    public init() {
+        self.creators = [:]
+    }
+
+    /// Creates a registry with given creators.
+    public init(creators: [String: @Sendable (URL) throws -> any LanguageModel]) {
+        self.creators = creators
+    }
+
+    /// Shared instance with default model types.
+    public static let shared: ModelTypeRegistry = .init(creators: all())
+
+    /// All predefined model types
+    private static func all() -> [String: @Sendable (URL) throws -> any LanguageModel] {
+        [
+            "paligemma": create(PaliGemmaConfiguration.self, PaliGemma.init),
+            "qwen2_vl": create(Qwen2VLConfiguration.self, Qwen2VL.init),
+            "idefics3": create(Idefics3Configuration.self, Idefics3.init),
+        ]
+    }
+
     // Note: using NSLock as we have very small (just dictionary get/set)
     // critical sections and expect no contention.  this allows the methods
     // to remain synchronous.
     private let lock = NSLock()
 
-    private var creators: [String: @Sendable (URL) throws -> any LanguageModel] = [
-        "paligemma": create(PaliGemmaConfiguration.self, PaliGemma.init),
-        "qwen2_vl": create(Qwen2VLConfiguration.self, Qwen2VL.init),
-        "qwen2_5_vl": create(Qwen25VLConfiguration.self, Qwen25VL.init),
-        "idefics3": create(Idefics3Configuration.self, Idefics3.init),
-    ]
+    private var creators: [String: @Sendable (URL) throws -> any LanguageModel]
 
     /// Add a new model to the type registry.
     public func registerModelType(
@@ -105,13 +122,25 @@ public class ModelTypeRegistry: @unchecked Sendable {
 
 public class ProcessorTypeRegistry: @unchecked Sendable {
 
-    // Note: using NSLock as we have very small (just dictionary get/set)
-    // critical sections and expect no contention.  this allows the methods
-    // to remain synchronous.
-    private let lock = NSLock()
+    /// Creates an empty registry.
+    public init() {
+        self.creators = [:]
+    }
 
-    private var creators:
-        [String: @Sendable (URL, any Tokenizer) throws -> any UserInputProcessor] = [
+    /// Creates a registry with given creators.
+    public init(creators: [String: @Sendable (URL, any Tokenizer) throws -> any UserInputProcessor])
+    {
+        self.creators = creators
+    }
+
+    /// Shared instance with default processor types.
+    public static let shared: ProcessorTypeRegistry = .init(creators: all())
+
+    /// All predefined processor types.
+    private static func all() -> [String: @Sendable (URL, any Tokenizer) throws ->
+        any UserInputProcessor]
+    {
+        [
             "PaliGemmaProcessor": create(
                 PaliGemmaProcessorConfiguration.self, PaliGemmaProcessor.init),
             "Qwen2VLProcessor": create(
@@ -121,6 +150,14 @@ public class ProcessorTypeRegistry: @unchecked Sendable {
             "Idefics3Processor": create(
                 Idefics3ProcessorConfiguration.self, Idefics3Processor.init),
         ]
+    }
+
+    // Note: using NSLock as we have very small (just dictionary get/set)
+    // critical sections and expect no contention.  this allows the methods
+    // to remain synchronous.
+    private let lock = NSLock()
+
+    private var creators: [String: @Sendable (URL, any Tokenizer) throws -> any UserInputProcessor]
 
     /// Add a new model to the type registry.
     public func registerProcessorType(
@@ -156,12 +193,21 @@ public class ProcessorTypeRegistry: @unchecked Sendable {
 /// swift-tokenizers code handles a good chunk of that and this is a place to augment that
 /// implementation, if needed.
 public class ModelRegistry: @unchecked Sendable {
+    /// Creates an empty registry.
+    public init() {
+        registry = Dictionary()
+    }
+
+    /// Creates a new registry with from given model configurations.
+    public init(modelConfigurations: [ModelConfiguration]) {
+        registry = Dictionary(uniqueKeysWithValues: modelConfigurations.map { ($0.name, $0) })
+    }
+
+    /// Shared instance with default model configurations.
+    public static let shared = ModelRegistry(modelConfigurations: all())
 
     private let lock = NSLock()
-    private var registry = Dictionary(
-        uniqueKeysWithValues: all().map {
-            ($0.name, $0)
-        })
+    private var registry: [String: ModelConfiguration]
 
     static public let paligemma3bMix448_8bit = ModelConfiguration(
         id: "mlx-community/paligemma-3b-mix-448-8bit",
@@ -188,6 +234,7 @@ public class ModelRegistry: @unchecked Sendable {
             paligemma3bMix448_8bit,
             qwen2VL2BInstruct4Bit,
             qwen2_5VL3BInstruct4Bit,
+            smolvlminstruct4bit,
         ]
     }
 
@@ -208,6 +255,12 @@ public class ModelRegistry: @unchecked Sendable {
             }
         }
     }
+
+    public var models: some Collection<ModelConfiguration> & Sendable {
+        lock.withLock {
+            return registry.values
+        }
+    }
 }
 
 /// Factory for creating new LLMs.
@@ -221,16 +274,27 @@ public class ModelRegistry: @unchecked Sendable {
 /// ```
 public class VLMModelFactory: ModelFactory {
 
-    public static let shared = VLMModelFactory()
+    public init(
+        typeRegistry: ModelTypeRegistry, processorRegistry: ProcessorTypeRegistry,
+        modelRegistry: ModelRegistry
+    ) {
+        self.typeRegistry = typeRegistry
+        self.processorRegistry = processorRegistry
+        self.modelRegistry = modelRegistry
+    }
+
+    /// Shared instance with default behavior.
+    public static let shared = VLMModelFactory(
+        typeRegistry: .shared, processorRegistry: .shared, modelRegistry: .shared)
 
     /// registry of model type, e.g. configuration value `paligemma` -> configuration and init methods
-    public let typeRegistry = ModelTypeRegistry()
+    public let typeRegistry: ModelTypeRegistry
 
     /// registry of input processor type, e.g. configuration value `PaliGemmaProcessor` -> configuration and init methods
-    public let processorRegistry = ProcessorTypeRegistry()
+    public let processorRegistry: ProcessorTypeRegistry
 
     /// registry of model id to configuration, e.g. `mlx-community/paligemma-3b-mix-448-8bit`
-    public let modelRegistry = ModelRegistry()
+    public let modelRegistry: ModelRegistry
 
     public func configuration(id: String) -> ModelConfiguration {
         modelRegistry.configuration(id: id)

Original file line number	Diff line number	Diff line change
`@@ -83,6 +83,12 @@ public struct ModelConfiguration: Sendable {`
`83`	`83`	`return ModelConfiguration(id: id)`
`84`	`84`	`}`
`85`	`85`	`}`
	`86`	`+`
	`87`	`+ @MainActor`
	`88`	`+ public static var models: some Collection<ModelConfiguration> & Sendable {`
	`89`	`+ bootstrap()`
	`90`	`+ return Self.registry.values`
	`91`	`+ }`
`86`	`92`	`}`
`87`	`93`
`88`	`94`	`extension ModelConfiguration {`