chore: import upstream snapshot with attribution

swift/StableDiffusion/tokenizer/BPETokenizer+Reading.swift

@@ -0,0 +1,58 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2022 Apple Inc. All Rights Reserved.
+
+import Foundation
+
+@available(iOS 16.2, macOS 13.1, *)
+extension BPETokenizer {
+    enum FileReadError: Error {
+        case invalidMergeFileLine(Int)
+    }
+
+    /// Read vocab.json file at URL into a dictionary mapping a String to its Int token id
+    static func readVocabulary(url: URL) throws -> [String: Int] {
+        let content = try Data(contentsOf: url)
+        return try JSONDecoder().decode([String: Int].self, from: content)
+    }
+
+    /// Read merges.txt file at URL into a dictionary mapping bigrams to the line number/rank/priority
+    static func readMerges(url: URL) throws -> [TokenPair: Int] {
+        let data = try Data(contentsOf: url)
+        var merges = [TokenPair: Int]()
+        var index = 0
+        var line = [UInt8]()
+        for byte in data {
+            if byte == UInt8(ascii: "\n") {
+                if let pair = try parseMergesLine(line, index: index) {
+                    merges[pair] = index
+                }
+                line.removeAll(keepingCapacity: true)
+                index += 1
+            } else {
+                line.append(byte)
+            }
+        }
+
+        return merges
+    }
+
+    static func parseMergesLine(_ line: [UInt8], index: Int) throws -> TokenPair? {
+        if line.isEmpty || line.first == UInt8(ascii: "#") {
+            return nil
+        }
+        let pair = line.split(separator: UInt8(ascii: " "))
+        if pair.count != 2 {
+            throw FileReadError.invalidMergeFileLine(index + 1)
+        }
+        return TokenPair( String(bytes: pair[0]), String(bytes: pair[1]))
+    }
+}
+
+extension String {
+    init(bytes: some Collection<UInt8>) {
+        self.init(unsafeUninitializedCapacity: bytes.count) { pointer in
+            _ = pointer.initialize(fromContentsOf: bytes)
+            return bytes.count
+        }
+    }
+}

swift/StableDiffusion/tokenizer/BPETokenizer.swift

@@ -0,0 +1,185 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2022 Apple Inc. All Rights Reserved.
+
+import Foundation
+
+/// A tokenizer based on byte pair encoding.
+@available(iOS 16.2, macOS 13.1, *)
+public struct BPETokenizer {
+    /// A dictionary that maps pairs of tokens to the rank/order of the merge.
+    let merges: [TokenPair : Int]
+
+    /// A dictionary from of tokens to identifiers.
+    let vocabulary: [String: Int]
+
+    /// The token used for padding
+    let padToken: String
+
+    /// The start token.
+    let startToken: String = "<|startoftext|>"
+
+    /// The end token.
+    let endToken: String = "<|endoftext|>"
+
+    /// The unknown token.
+    let unknownToken: String = "<|endoftext|>"
+
+    var unknownTokenID: Int {
+        vocabulary[unknownToken, default: 0]
+    }
+
+    /// Creates a tokenizer.
+    ///
+    /// - Parameters:
+    ///   - merges: A dictionary that maps pairs of tokens to the rank/order of the merge.
+    ///   - vocabulary: A dictionary from of tokens to identifiers.
+    public init(merges: [TokenPair: Int], vocabulary: [String: Int], padToken: String = "<|endoftext|>") {
+        self.merges = merges
+        self.vocabulary = vocabulary
+        self.padToken = padToken
+    }
+
+    /// Creates a tokenizer by loading merges and vocabulary from URLs.
+    ///
+    /// - Parameters:
+    ///   - mergesURL: The URL of a text file containing merges.
+    ///   - vocabularyURL: The URL of a JSON file containing the vocabulary.
+    public init(mergesAt mergesURL: URL, vocabularyAt vocabularyURL: URL, padToken: String = "<|endoftext|>") throws {
+        self.merges = try Self.readMerges(url: mergesURL)
+        self.vocabulary = try! Self.readVocabulary(url: vocabularyURL)
+        self.padToken = padToken
+    }
+
+    /// Tokenizes an input string.
+    ///
+    /// - Parameters:
+    ///   - input: A string.
+    ///   - minCount: The minimum number of tokens to return.
+    /// - Returns: An array of tokens and an array of token identifiers.
+    public func tokenize(input: String, minCount: Int? = nil) -> (tokens: [String], tokenIDs: [Int]) {
+        var tokens: [String] = []
+
+        tokens.append(startToken)
+        tokens.append(contentsOf: encode(input: input))
+        tokens.append(endToken)
+
+        // Pad if there was a min length specified
+        if let minLen = minCount, minLen > tokens.count {
+            tokens.append(contentsOf: repeatElement(padToken, count: minLen - tokens.count))
+        }
+
+        let ids = tokens.map({ vocabulary[$0, default: unknownTokenID] })
+        return (tokens: tokens, tokenIDs: ids)
+    }
+
+    /// Returns the token identifier for a token.
+    public func tokenID(for token: String) -> Int? {
+        vocabulary[token]
+    }
+
+    /// Returns the token for a token identifier.
+    public func token(id: Int) -> String? {
+        vocabulary.first(where: { $0.value == id })?.key
+    }
+
+    /// Decodes a sequence of tokens into a fully formed string
+    public func decode(tokens: [String]) -> String {
+        String(tokens.joined())
+            .replacingOccurrences(of: "</w>", with: " ")
+            .replacingOccurrences(of: startToken, with: "")
+            .replacingOccurrences(of: endToken, with: "")
+    }
+
+    /// Encode an input string to a sequence of tokens
+    func encode(input: String) -> [String] {
+        let normalized = input.trimmingCharacters(in: .whitespacesAndNewlines).lowercased()
+        let words = normalized.split(separator: " ")
+        return words.flatMap({ encode(word: $0) })
+    }
+
+    /// Encode a single word into a sequence of tokens
+    func encode(word: Substring) -> [String] {
+        var tokens = word.map { String($0) }
+        if let last = tokens.indices.last {
+            tokens[last] = tokens[last] + "</w>"
+        }
+
+        while true {
+            let pairs = pairs(for: tokens)
+            let canMerge = pairs.filter { merges[$0] != nil }
+
+            if canMerge.isEmpty {
+                break
+            }
+
+            // If multiple merges are found, use the one with the lowest rank
+            let shouldMerge = canMerge.min { merges[$0]! < merges[$1]! }!
+            tokens = update(tokens, merging: shouldMerge)
+        }
+        return tokens
+    }
+
+    /// Get  the set of adjacent pairs / bigrams from a sequence of tokens
+    func pairs(for tokens: [String]) -> Set<TokenPair> {
+        guard tokens.count > 1 else {
+            return Set()
+        }
+
+        var pairs = Set<TokenPair>(minimumCapacity: tokens.count - 1)
+        var prev = tokens.first!
+        for current in tokens.dropFirst() {
+            pairs.insert(TokenPair(prev, current))
+            prev = current
+        }
+        return pairs
+    }
+
+    /// Update the sequence of tokens by greedily merging instance of a specific bigram
+    func update(_ tokens: [String], merging bigram: TokenPair) -> [String] {
+        guard tokens.count > 1 else {
+            return []
+        }
+
+        var newTokens = [String]()
+        newTokens.reserveCapacity(tokens.count - 1)
+
+        var index = 0
+        while index < tokens.count {
+            let remainingTokens = tokens[index...]
+            if let startMatchIndex = remainingTokens.firstIndex(of: bigram.first) {
+                // Found a possible match, append everything before it
+                newTokens.append(contentsOf: tokens[index..<startMatchIndex])
+
+                if index < tokens.count - 1 && tokens[startMatchIndex + 1] == bigram.second {
+                    // Full match, merge
+                    newTokens.append(bigram.first + bigram.second)
+                    index = startMatchIndex + 2
+                } else {
+                    // Only matched the first, no merge
+                    newTokens.append(bigram.first)
+                    index = startMatchIndex + 1
+                }
+            } else {
+                // Didn't find any more matches, append the rest unmerged
+                newTokens.append(contentsOf: remainingTokens)
+                break
+            }
+        }
+        return newTokens
+    }
+}
+
+@available(iOS 16.2, macOS 13.1, *)
+extension BPETokenizer {
+
+    /// A hashable tuple of strings
+    public struct TokenPair: Hashable {
+        let first: String
+        let second: String
+
+        init(_ first: String, _ second: String) {
+            self.first = first
+            self.second = second
+        }
+    }
+}

swift/StableDiffusion/tokenizer/T5Tokenizer.swift

@@ -0,0 +1,21 @@
+// For licensing see accompanying LICENSE.md file.
+// Copyright (C) 2024 Apple Inc. All Rights Reserved.
+
+import Foundation
+import Hub
+import Tokenizers
+
+/// Extension to swift-transfomers Hub.swift to load local Config files
+public extension Config {
+    /// Assumes the file is already present at local url.
+    /// `fileURL` is a complete local file path for the given model
+    public init(fileURL: URL) throws  {
+        let data = try Data(contentsOf: fileURL)
+        let parsed = try JSONSerialization.jsonObject(with: data, options: [])
+        guard var dictionary = parsed as? [String: Any] else { throw Hub.HubClientError.parse }
+        
+        // Necessary override for loading local tokenizer configs
+        dictionary["tokenizer_class"] = "T5Tokenizer"
+        self.init(dictionary)
+    }
+}