3928858 · splitfireai/demucs-swift

@@ -0,0 +1,381 @@

1 +import Foundation

2 +

3 +/// One source the models can separate out of a mix.

4 +public enum StemKind: String, Sendable, CaseIterable {

5 + case drums

6 + case bass

7 + case other

8 + case vocals

9 + case guitar

10 + case piano

11 +}

12 +

13 +/// Which HTDemucs variant to download or load.

14 +public enum DemucsModel: Sendable, Hashable {

15 + /// htdemucs: drums, bass, other, vocals. ~84 MB.

16 + case fourStem

17 +

18 + /// htdemucs_6s: adds guitar and piano. ~84 MB.

19 + case sixStem

20 +

21 + /// htdemucs_ft: one fine-tuned sub-model per stem, best quality. ~333 MB.

22 + ///

23 + /// Only drums, bass, other and vocals are available. Each requested stem

24 + /// runs its own sub-model, so asking for fewer is proportionally faster.

25 + /// Every selection shares the same downloaded weights.

26 + case fineTuned(stems: [StemKind])

27 +

28 + /// The fine-tuned model with all four of its stems.

29 + public static let fineTunedAll = DemucsModel.fineTuned(stems: [.drums, .bass, .other, .vocals])

30 +}

31 +

32 +public struct DemucsModelInfo: Sendable, Hashable, Identifiable {

33 + public let model: DemucsModel

34 + /// The model's upstream name, such as `htdemucs`.

35 + public let id: String

36 + public let label: String

37 + public let description: String

38 + public let sizeMB: UInt32

39 + public let stems: [StemKind]

40 +}

41 +

42 +/// A separated source at the input sample rate, the same length as the input.

43 +public struct Stem: Sendable, Equatable {

44 + public let kind: StemKind

45 + public let left: [Float]

46 + public let right: [Float]

47 +}

48 +

49 +public struct DownloadProgress: Sendable, Equatable {

50 + public let downloadedBytes: UInt64

51 + /// From the server's `Content-Length`, or the model's nominal size.

52 + public let totalBytes: UInt64

53 + public let fraction: Double

54 +}

55 +

56 +public struct SeparationProgress: Sendable, Equatable {

57 + public let fraction: Double

58 + public let chunkIndex: UInt64

59 + public let totalChunks: UInt64

60 +}

61 +

62 +public enum DemucsError: Error, Sendable, Equatable, LocalizedError {

63 + /// The task was cancelled, or a progress handler asked to stop.

64 + case cancelled

65 + /// `separate` or `warmup` was called before `load`.

66 + case notLoaded

67 + /// `load` was called before the model's weights were downloaded.

68 + case notDownloaded(modelID: String)

69 + case invalidInput(String)

70 + case downloadFailed(String)

71 + case io(String)

72 + case inferenceFailed(String)

73 +

74 + public var errorDescription: String? {

75 + switch self {

76 + case .cancelled: "The operation was cancelled."

77 + case .notLoaded: "No model is loaded."

78 + case .notDownloaded(let id): "Model \(id) has not been downloaded."

79 + case .invalidInput(let reason): reason

80 + case .downloadFailed(let reason): "Download failed: \(reason)"

81 + case .io(let reason): "File error: \(reason)"

82 + case .inferenceFailed(let reason): reason

83 + }

84 + }

85 +}

86 +

87 +/// On-device music source separation (HTDemucs v4), GPU-accelerated through

88 +/// Metal.

89 +///

90 +/// An engine holds at most one loaded model. Loading, warming up and

91 +/// separating on one engine run one at a time; create another engine to run

92 +/// a second model side by side, memory permitting.

93 +///

94 +/// ```swift

95 +/// try await DemucsEngine.download(.fourStem)

96 +/// let engine = DemucsEngine()

97 +/// try await engine.load(.fourStem)

98 +/// let stems = try await engine.separate(left: left, right: right, sampleRate: 44_100)

99 +/// ```

100 +public actor DemucsEngine {

101 + private nonisolated let core: FfiDemucsEngine

102 +

103 + public init() {

104 + core = FfiDemucsEngine()

105 + }

106 +

107 + // MARK: Model files

108 +

109 + /// Sets where model weights are stored. Call once at launch, before any

110 + /// other Demucs call; sandboxed apps should pass a directory inside their

111 + /// container, such as one under Caches. Without it, weights go to the

112 + /// platform cache directory under `demucs-rs/`.

113 + public static func configureCacheDirectory(_ url: URL) {

114 + configureCacheDir(path: url.path)

115 + }

116 +

117 + public static var models: [DemucsModelInfo] {

118 + allModels().map(\.swift)

119 + }

120 +

121 + public static func info(for model: DemucsModel) -> DemucsModelInfo {

122 + modelMetadata(model: model.ffi).swift

123 + }

124 +

125 + public static func isDownloaded(_ model: DemucsModel) -> Bool {

126 + isModelCached(model: model.ffi)

127 + }

128 +

129 + /// The downloaded weights file, or `nil` if it is not on disk.

130 + public static func downloadedFileURL(for model: DemucsModel) -> URL? {

131 + cachedModelPath(model: model.ffi).map { URL(fileURLWithPath: $0) }

132 + }

133 +

134 + /// Removes the downloaded weights. Returns `false` if there was nothing

135 + /// to remove.

136 + @discardableResult

137 + public static func deleteDownload(_ model: DemucsModel) -> Bool {

138 + deleteCachedModel(model: model.ffi)

139 + }

140 +

141 + /// Downloads the model's weights from Hugging Face. Returns straight away

142 + /// if they are already downloaded. Cancel the calling task to stop; a

143 + /// partial download is discarded.

144 + ///

145 + /// `progress` is called on a background thread.

146 + public static func download(

147 + _ model: DemucsModel,

148 + progress: (@Sendable (DownloadProgress) -> Void)? = nil

149 + ) async throws {

150 + let listener = DownloadListener(handler: progress)

151 + try await withTaskCancellationHandler {

152 + try await translatingErrors {

153 + try await downloadModel(model: model.ffi, listener: listener)

154 + }

155 + } onCancel: {

156 + listener.cancel()

157 + }

158 + }

159 +

160 + // MARK: Loading

161 +

162 + /// Loads a downloaded model, replacing the current one. Returns the load

163 + /// time in seconds.

164 + @discardableResult

165 + public func load(_ model: DemucsModel) async throws -> Double {

166 + try await translatingErrors {

167 + try await core.loadModel(model: model.ffi)

168 + }

169 + }

170 +

171 + /// Runs a dummy separation so the first real one does not pay for GPU

172 + /// shader compilation. Worth calling once after `load`, off the critical

173 + /// path.

174 + public func warmup() async throws {

175 + try await translatingErrors {

176 + try await core.warmup()

177 + }

178 + }

179 +

180 + /// Frees the loaded model. Returns `false` if nothing was loaded.

181 + @discardableResult

182 + public func unload() async -> Bool {

183 + await core.unloadModel()

184 + }

185 +

186 + public nonisolated var isLoaded: Bool {

187 + core.isLoaded()

188 + }

189 +

190 + public nonisolated var loadedModel: DemucsModel? {

191 + core.loadedModel()?.swift

192 + }

193 +

194 + // MARK: Separation

195 +

196 + /// Separates stereo PCM at any sample rate into the loaded model's stems.

197 + ///

198 + /// Both channels must have the same, non-zero length. Audio is processed

199 + /// in chunks of about 7.8 seconds; cancelling the calling task stops at

200 + /// the next chunk boundary. `progress` is called on a background thread.

201 + public func separate(

202 + left: [Float],

203 + right: [Float],

204 + sampleRate: UInt32,

205 + progress: (@Sendable (SeparationProgress) -> Void)? = nil

206 + ) async throws -> [Stem] {

207 + let listener = SeparationListener(handler: progress)

208 + let stems = try await withTaskCancellationHandler {

209 + try await translatingErrors {

210 + try await core.separate(

211 + left: left,

212 + right: right,

213 + sampleRate: sampleRate,

214 + listener: listener

215 + )

216 + }

217 + } onCancel: {

218 + listener.cancel()

219 + }

220 + return stems.map(\.swift)

221 + }

222 +

223 + /// Separates mono PCM. The stems come back stereo with identical channels.

224 + public func separate(

225 + mono samples: [Float],

226 + sampleRate: UInt32,

227 + progress: (@Sendable (SeparationProgress) -> Void)? = nil

228 + ) async throws -> [Stem] {

229 + try await separate(left: samples, right: samples, sampleRate: sampleRate, progress: progress)

230 + }

231 +

232 + /// How many chunks `separate` splits this many samples into, for sizing

233 + /// progress UI up front.

234 + public static func chunkCount(sampleCount: Int, sampleRate: UInt32) -> Int {

235 + Int(numChunks(nSamples: UInt64(max(sampleCount, 0)), sampleRate: sampleRate))

236 + }

237 +}

238 +

239 +// MARK: - Callback adapters

240 +

241 +/// Shared by both listeners: the progress callback returns `false` once the

242 +/// Swift task is cancelled, which is how Rust learns to stop.

243 +private final class CancellationFlag: @unchecked Sendable {

244 + private let lock = NSLock()

245 + private var cancelled = false

246 +

247 + var isCancelled: Bool { lock.withLock { cancelled } }

248 +

249 + func cancel() { lock.withLock { cancelled = true } }

250 +}

251 +

252 +private final class DownloadListener: FfiDownloadListener, @unchecked Sendable {

253 + private let handler: (@Sendable (DownloadProgress) -> Void)?

254 + private let flag = CancellationFlag()

255 +

256 + init(handler: (@Sendable (DownloadProgress) -> Void)?) {

257 + self.handler = handler

258 + }

259 +

260 + func cancel() { flag.cancel() }

261 +

262 + func onProgress(progress: FfiDownloadProgress) -> Bool {

263 + guard !flag.isCancelled else { return false }

264 + handler?(DownloadProgress(

265 + downloadedBytes: progress.downloadedBytes,

266 + totalBytes: progress.totalBytes,

267 + fraction: progress.fraction

268 + ))

269 + return true

270 + }

271 +}

272 +

273 +private final class SeparationListener: FfiSeparationListener, @unchecked Sendable {

274 + private let handler: (@Sendable (SeparationProgress) -> Void)?

275 + private let flag = CancellationFlag()

276 +

277 + init(handler: (@Sendable (SeparationProgress) -> Void)?) {

278 + self.handler = handler

279 + }

280 +

281 + func cancel() { flag.cancel() }

282 +

283 + func onProgress(progress: FfiSeparationProgress) -> Bool {

284 + guard !flag.isCancelled else { return false }

285 + handler?(SeparationProgress(

286 + fraction: progress.fraction,

287 + chunkIndex: progress.chunkIndex,

288 + totalChunks: progress.totalChunks

289 + ))

290 + return true

291 + }

292 +}

293 +

294 +// MARK: - Conversions

295 +

296 +private func translatingErrors<T>(_ body: () async throws -> T) async throws -> T {

297 + do {

298 + return try await body()

299 + } catch let error as FfiDemucsError {

300 + throw DemucsError(error)

301 + }

302 +}

303 +

304 +extension DemucsError {

305 + init(_ error: FfiDemucsError) {

306 + switch error {

307 + case .Cancelled: self = .cancelled

308 + case .NotLoaded: self = .notLoaded

309 + case .NotCached(let modelId): self = .notDownloaded(modelID: modelId)

310 + case .InvalidInput(let reason): self = .invalidInput(reason)

311 + case .Download(let reason): self = .downloadFailed(reason)

312 + case .Io(let reason): self = .io(reason)

313 + case .Inference(let reason): self = .inferenceFailed(reason)

314 + }

315 + }

316 +}

317 +

318 +extension StemKind {

319 + var ffi: FfiStemKind {

320 + switch self {

321 + case .drums: .drums

322 + case .bass: .bass

323 + case .other: .other

324 + case .vocals: .vocals

325 + case .guitar: .guitar

326 + case .piano: .piano

327 + }

328 + }

329 +}

330 +

331 +extension FfiStemKind {

332 + var swift: StemKind {

333 + switch self {

334 + case .drums: .drums

335 + case .bass: .bass

336 + case .other: .other

337 + case .vocals: .vocals

338 + case .guitar: .guitar

339 + case .piano: .piano

340 + }

341 + }

342 +}

343 +

344 +extension DemucsModel {

345 + var ffi: FfiDemucsModel {

346 + switch self {

347 + case .fourStem: .fourStem

348 + case .sixStem: .sixStem

349 + case .fineTuned(let stems): .fineTuned(stems: stems.map(\.ffi))

350 + }

351 + }

352 +}

353 +

354 +extension FfiDemucsModel {

355 + var swift: DemucsModel {

356 + switch self {

357 + case .fourStem: .fourStem

358 + case .sixStem: .sixStem

359 + case .fineTuned(let stems): .fineTuned(stems: stems.map(\.swift))

360 + }

361 + }

362 +}

363 +

364 +private extension FfiModelMetadata {

365 + var swift: DemucsModelInfo {

366 + DemucsModelInfo(

367 + model: model.swift,

368 + id: id,

369 + label: label,

370 + description: description,

371 + sizeMB: sizeMb,

372 + stems: stems.map(\.swift)

373 + )

374 + }

375 +}

376 +

377 +private extension FfiStem {

378 + var swift: Stem {

379 + Stem(kind: kind.swift, left: left, right: right)

380 + }

381 +}

Add a Swift API over the generated bindings