Add Xet support for faster downloads

Author: mattt

Package.swift

@@ -20,13 +20,15 @@ let package = Package(
         )
     ],
     dependencies: [
-        .package(url: "https://github.com/mattt/EventSource.git", from: "1.0.0")
+        .package(url: "https://github.com/mattt/EventSource.git", from: "1.0.0"),
+        .package(url: "https://github.com/mattt/swift-xet.git", branch: "main"),
     ],
     targets: [
         .target(
             name: "HuggingFace",
             dependencies: [
-                .product(name: "EventSource", package: "EventSource")
+                .product(name: "EventSource", package: "EventSource"),
+                .product(name: "Xet", package: "swift-xet"),
             ],
             path: "Sources/HuggingFace"
         ),

Sources/HuggingFace/Hub/HubClient+Files.swift

@@ -6,6 +6,10 @@ import UniformTypeIdentifiers
     import FoundationNetworking
 #endif
 
+#if canImport(Xet)
+    import Xet
+#endif
+
 // MARK: - Upload Operations
 
 public extension HubClient {
@@ -183,6 +187,33 @@ public extension HubClient {
         useRaw: Bool = false,
         cachePolicy: URLRequest.CachePolicy = .useProtocolCachePolicy
     ) async throws -> Data {
+        #if canImport(Xet)
+            if isXetEnabled {
+                do {
+                    let tempDirectory = FileManager.default.temporaryDirectory
+                        .appendingPathComponent(UUID().uuidString, isDirectory: true)
+                    let tempFile = tempDirectory.appendingPathComponent(UUID().uuidString)
+                    try FileManager.default.createDirectory(at: tempDirectory, withIntermediateDirectories: true)
+                    defer { try? FileManager.default.removeItem(at: tempDirectory) }
+
+                    if try await downloadFileWithXet(
+                        repoPath: repoPath,
+                        repo: repo,
+                        revision: revision,
+                        destination: tempFile,
+                        progress: nil
+                    ) != nil {
+                        return try Data(contentsOf: tempFile)
+                    } else {
+                        print("⚠️  Xet returned nil for \(repoPath), falling back to LFS")
+                    }
+                } catch {
+                    print("⚠️  Xet failed for \(repoPath): \(error), falling back to LFS")
+                }
+            }
+        #endif
+
+        // Fallback to existing LFS download method
         let endpoint = useRaw ? "raw" : "resolve"
         let urlPath = "/\(repo)/\(endpoint)/\(revision)/\(repoPath)"
         var request = try httpClient.createRequest(.get, urlPath)
@@ -215,6 +246,27 @@ public extension HubClient {
         cachePolicy: URLRequest.CachePolicy = .useProtocolCachePolicy,
         progress: Progress? = nil
     ) async throws -> URL {
+        #if canImport(Xet)
+            if isXetEnabled {
+                do {
+                    if let downloaded = try await downloadFileWithXet(
+                        repoPath: repoPath,
+                        repo: repo,
+                        revision: revision,
+                        destination: destination,
+                        progress: progress
+                    ) {
+                        return downloaded
+                    } else {
+                        print("⚠️  Xet returned nil for \(repoPath), falling back to LFS")
+                    }
+                } catch {
+                    print("⚠️  Xet failed for \(repoPath): \(error), falling back to LFS")
+                }
+            }
+        #endif
+
+        // Fallback to existing LFS download method
         let endpoint = useRaw ? "raw" : "resolve"
         let urlPath = "/\(repo)/\(endpoint)/\(revision)/\(repoPath)"
         var request = try httpClient.createRequest(.get, urlPath)
@@ -542,6 +594,81 @@ public extension HubClient {
     }
 }
 
+#if canImport(Xet)
+    private extension HubClient {
+        /// Downloads a file using Xet's content-addressable storage system.
+        ///
+        /// This method uses a cached XetClient instance and JWT tokens to maximize
+        /// download performance through connection reuse and reduced API overhead.
+        ///
+        /// Performance optimizations:
+        /// - Reuses a single XetClient across all downloads for HTTP/TLS connection pooling
+        /// - Caches CAS JWT tokens per (repo, revision) to avoid redundant API calls
+        /// - Leverages Xet's parallel chunk downloading (configurable via XET_MAX_PARALLEL_FILES env var)
+        @discardableResult
+        func downloadFileWithXet(
+            repoPath: String,
+            repo: Repo.ID,
+            revision: String,
+            destination: URL,
+            progress: Progress?
+        ) async throws -> URL? {
+            guard isXetEnabled else {
+                return nil
+            }
+
+            let xetClient = try getXetClient()
+
+            guard
+                let fileInfo = try xetClient.getFileInfo(
+                    repo: repo.rawValue,
+                    path: repoPath,
+                    revision: revision
+                )
+            else {
+                return nil
+            }
+
+            let jwt = try getCachedJwt(
+                xetClient: xetClient,
+                repo: repo.rawValue,
+                revision: revision,
+                isUpload: false
+            )
+
+            let destinationDirectory = destination.deletingLastPathComponent()
+            try FileManager.default.createDirectory(
+                at: destinationDirectory,
+                withIntermediateDirectories: true
+            )
+
+            let downloads = try xetClient.downloadFiles(
+                fileInfos: [fileInfo],
+                destinationDir: destinationDirectory.path,
+                jwtInfo: jwt
+            )
+
+            guard let downloadedPath = downloads.first else {
+                return nil
+            }
+
+            let downloadedURL = URL(fileURLWithPath: downloadedPath)
+
+            if downloadedURL.standardizedFileURL != destination.standardizedFileURL {
+                if FileManager.default.fileExists(atPath: destination.path) {
+                    try FileManager.default.removeItem(at: destination)
+                }
+                try FileManager.default.moveItem(at: downloadedURL, to: destination)
+            }
+
+            progress?.totalUnitCount = 100
+            progress?.completedUnitCount = 100
+
+            return destination
+        }
+    }
+#endif
+
 // MARK: - Metadata Helpers
 
 extension HubClient {

Sources/HuggingFace/Hub/HubClient.swift

@@ -4,6 +4,10 @@ import Foundation
     import FoundationNetworking
 #endif
 
+#if canImport(Xet)
+    import Xet
+#endif
+
 /// A Hugging Face Hub API client.
 ///
 /// This client provides methods to interact with the Hugging Face Hub API,
@@ -32,8 +36,19 @@ public final class HubClient: Sendable {
     /// environment variable (defaults to https://huggingface.co).
     public static let `default` = HubClient()
 
+    /// Indicates whether Xet acceleration is enabled for this client.
+    public let isXetEnabled: Bool
+
     /// The underlying HTTP client.
     internal let httpClient: HTTPClient
+    
+    #if canImport(Xet)
+        /// Xet client instance for connection reuse (created once during initialization)
+        private let xetClient: XetClient?
+        
+        /// Thread-safe JWT cache for CAS access tokens
+        private let jwtCache: JwtCache
+    #endif
 
     /// The host URL for requests made by the client.
     public var host: URL {
@@ -67,13 +82,15 @@ public final class HubClient: Sendable {
     ///   - userAgent: The value for the `User-Agent` header sent in requests, if any. Defaults to `nil`.
     public convenience init(
         session: URLSession = URLSession(configuration: .default),
-        userAgent: String? = nil
+        userAgent: String? = nil,
+        enableXet: Bool = HubClient.isXetSupported
     ) {
         self.init(
             session: session,
             host: Self.detectHost(),
             userAgent: userAgent,
-            tokenProvider: .environment
+            tokenProvider: .environment,
+            enableXet: enableXet
         )
     }
 
@@ -88,13 +105,15 @@ public final class HubClient: Sendable {
         session: URLSession = URLSession(configuration: .default),
         host: URL,
         userAgent: String? = nil,
-        bearerToken: String? = nil
+        bearerToken: String? = nil,
+        enableXet: Bool = HubClient.isXetSupported
     ) {
         self.init(
             session: session,
             host: host,
             userAgent: userAgent,
-            tokenProvider: bearerToken.map { .fixed(token: $0) } ?? .none
+            tokenProvider: bearerToken.map { .fixed(token: $0) } ?? .none,
+            enableXet: enableXet
         )
     }
 
@@ -109,14 +128,28 @@ public final class HubClient: Sendable {
         session: URLSession = URLSession(configuration: .default),
         host: URL,
         userAgent: String? = nil,
-        tokenProvider: TokenProvider
+        tokenProvider: TokenProvider,
+        enableXet: Bool = HubClient.isXetSupported
     ) {
+        self.isXetEnabled = enableXet && HubClient.isXetSupported
         self.httpClient = HTTPClient(
             host: host,
             userAgent: userAgent,
             tokenProvider: tokenProvider,
             session: session
         )
+        
+        #if canImport(Xet)
+            self.jwtCache = JwtCache()
+            
+            if self.isXetEnabled {
+                // Create XetClient once during initialization
+                let token = try? tokenProvider.getToken()
+                self.xetClient = try? (token.map { try XetClient.withToken(token: $0) } ?? XetClient())
+            } else {
+                self.xetClient = nil
+            }
+        #endif
     }
 
     // MARK: - Auto-detection
@@ -134,4 +167,104 @@ public final class HubClient: Sendable {
         }
         return defaultHost
     }
-}
+
+    public static var isXetSupported: Bool {
+        #if canImport(Xet)
+            return true
+        #else
+            return false
+        #endif
+    }
+
+    // MARK: - Xet Client
+
+    #if canImport(Xet)
+        /// Thread-safe cache for CAS JWT tokens
+        private final class JwtCache: @unchecked Sendable {
+            private struct CacheKey: Hashable {
+                let repo: String
+                let revision: String
+            }
+            
+            private struct CachedJwt {
+                let jwt: CasJwtInfo
+                let expiresAt: Date
+                
+                var isExpired: Bool {
+                    Date() >= expiresAt
+                }
+            }
+            
+            private var cache: [CacheKey: CachedJwt] = [:]
+            private let lock = NSLock()
+            
+            func get(repo: String, revision: String) -> CasJwtInfo? {
+                lock.lock()
+                defer { lock.unlock() }
+                
+                let key = CacheKey(repo: repo, revision: revision)
+                if let cached = cache[key], !cached.isExpired {
+                    return cached.jwt
+                }
+                return nil
+            }
+            
+            func set(jwt: CasJwtInfo, repo: String, revision: String) {
+                lock.lock()
+                defer { lock.unlock() }
+                
+                let key = CacheKey(repo: repo, revision: revision)
+                // Cache with expiration (5 minutes before actual expiry for safety)
+                let expiresAt = Date(timeIntervalSince1970: TimeInterval(jwt.exp())) - 300
+                cache[key] = CachedJwt(jwt: jwt, expiresAt: expiresAt)
+            }
+        }
+    
+        /// Returns the Xet client for faster downloads.
+        ///
+        /// The client is created once during initialization and reused across downloads
+        /// to enable connection pooling and avoid reinitialization overhead.
+        ///
+        /// - Returns: A Xet client instance.
+        internal func getXetClient() throws -> XetClient {
+            guard isXetEnabled, let client = xetClient else {
+                throw HTTPClientError.requestError("Xet support is disabled for this client.")
+            }
+            return client
+        }
+        
+        /// Gets or fetches a CAS JWT for the given repository and revision.
+        ///
+        /// JWTs are cached to avoid redundant API calls.
+        ///
+        /// - Parameters:
+        ///   - xetClient: The Xet client to use for fetching the JWT
+        ///   - repo: Repository identifier
+        ///   - revision: Git revision
+        ///   - isUpload: Whether this JWT is for upload (true) or download (false)
+        /// - Returns: A CAS JWT info object
+        internal func getCachedJwt(
+            xetClient: XetClient,
+            repo: String,
+            revision: String,
+            isUpload: Bool
+        ) throws -> CasJwtInfo {
+            // Check cache first
+            if let cached = jwtCache.get(repo: repo, revision: revision) {
+                return cached
+            }
+            
+            // Fetch a new JWT
+            let jwt = try xetClient.getCasJwt(
+                repo: repo,
+                revision: revision,
+                isUpload: isUpload
+            )
+            
+            // Cache it
+            jwtCache.set(jwt: jwt, repo: repo, revision: revision)
+            
+            return jwt
+        }
+    #endif
+    }

Tests/HuggingFaceTests/HubTests/HubClientTests.swift

@@ -38,4 +38,17 @@ struct HubClientTests {
 
         #expect(client.host.path.hasSuffix("/"))
     }
+
+    @Test("Xet configuration can be toggled per client")
+    func testXetConfigurationToggle() throws {
+        let host = URL(string: "https://huggingface.co")!
+
+        try #require(HubClient.isXetSupported, "Xet is not supported on this platform")
+
+        let disabledClient = HubClient(host: host, enableXet: false)
+        #expect(disabledClient.isXetEnabled == false)
+
+        let enabledClient = HubClient(host: host, enableXet: true)
+        #expect(enabledClient.isXetEnabled)
+    }
 }