@splitfireai / demucs-swift
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
swift 167 lines 6.12 KB
Raw
1 import AVFoundation
2 import Demucs
3 import Foundation
4
5 struct StemFile: Identifiable {
6 let kind: StemKind
7 let url: URL
8 var id: StemKind { kind }
9 }
10
11 @MainActor
12 final class SeparationViewModel: ObservableObject {
13 enum Phase: Equatable {
14 case idle
15 case downloading(Double)
16 case loading
17 case separating(Double)
18 case done
19 case failed(String)
20 }
21
22 @Published var model: DemucsModel = .fourStem
23 @Published private(set) var phase: Phase = .idle
24 @Published private(set) var stems: [StemFile] = []
25
26 private let engine = DemucsEngine()
27 private var task: Task<Void, Never>?
28 /// Bumped on every new run so a cancelled run that unwinds late can't
29 /// overwrite the state of the run that replaced it.
30 private var generation = 0
31
32 var isBusy: Bool {
33 switch phase {
34 case .downloading, .loading, .separating: true
35 default: false
36 }
37 }
38
39 func separate(fileAt url: URL) {
40 task?.cancel()
41 stems = []
42 generation += 1
43 let current = generation
44 task = Task { await run(url: url, generation: current) }
45 }
46
47 func cancel() {
48 task?.cancel()
49 }
50
51 /// Progress arrives from background threads and can land after the step
52 /// it describes has finished, or after a newer run has started; only
53 /// apply it while that run and step are current.
54 private func update(_ progress: Phase, generation current: Int) {
55 guard current == generation else { return }
56 switch (phase, progress) {
57 case (.downloading, .downloading), (.separating, .separating):
58 phase = progress
59 default:
60 break
61 }
62 }
63
64 private func run(url: URL, generation current: Int) async {
65 do {
66 let model = model
67 if !DemucsEngine.isDownloaded(model) {
68 phase = .downloading(0)
69 try await DemucsEngine.download(model) { progress in
70 Task { @MainActor in self.update(.downloading(progress.fraction), generation: current) }
71 }
72 }
73 if engine.loadedModel != model {
74 phase = .loading
75 try await engine.load(model)
76 }
77
78 let audio = try AudioFile.read(url)
79 phase = .separating(0)
80 let separated = try await engine.separate(
81 left: audio.left,
82 right: audio.right,
83 sampleRate: audio.sampleRate
84 ) { progress in
85 Task { @MainActor in self.update(.separating(progress.fraction), generation: current) }
86 }
87
88 let outputDirectory = FileManager.default.temporaryDirectory
89 .appendingPathComponent(url.deletingPathExtension().lastPathComponent, isDirectory: true)
90 try FileManager.default.createDirectory(at: outputDirectory, withIntermediateDirectories: true)
91 guard current == generation else { return }
92 stems = try separated.map { stem in
93 let file = outputDirectory.appendingPathComponent("\(stem.kind.rawValue).wav")
94 try AudioFile.write(stem, sampleRate: audio.sampleRate, to: file)
95 return StemFile(kind: stem.kind, url: file)
96 }
97 guard current == generation else { return }
98 phase = .done
99 } catch {
100 guard current == generation else { return }
101 switch error {
102 case DemucsError.cancelled, is CancellationError:
103 phase = .idle
104 default:
105 phase = .failed(error.localizedDescription)
106 }
107 }
108 }
109 }
110
111 /// Minimal AVFoundation decode/encode. Demucs itself takes and returns plain
112 /// per-channel Float arrays at any sample rate.
113 enum AudioFile {
114 struct Decoded {
115 let left: [Float]
116 let right: [Float]
117 let sampleRate: UInt32
118 }
119
120 static func read(_ url: URL) throws -> Decoded {
121 let scoped = url.startAccessingSecurityScopedResource()
122 defer { if scoped { url.stopAccessingSecurityScopedResource() } }
123
124 // Force Float32 so 16/24-bit PCM files decode into floatChannelData.
125 let file = try AVAudioFile(forReading: url, commonFormat: .pcmFormatFloat32, interleaved: false)
126 let format = file.processingFormat
127 guard let buffer = AVAudioPCMBuffer(pcmFormat: format, frameCapacity: AVAudioFrameCount(file.length)) else {
128 throw DemucsError.invalidInput("Could not allocate an audio buffer.")
129 }
130 try file.read(into: buffer)
131 guard let channels = buffer.floatChannelData, buffer.frameLength > 0 else {
132 throw DemucsError.invalidInput("The file has no audio.")
133 }
134 let frames = Int(buffer.frameLength)
135 let left = Array(UnsafeBufferPointer(start: channels[0], count: frames))
136 let right = format.channelCount > 1
137 ? Array(UnsafeBufferPointer(start: channels[1], count: frames))
138 : left
139 return Decoded(left: left, right: right, sampleRate: UInt32(format.sampleRate))
140 }
141
142 static func write(_ stem: Stem, sampleRate: UInt32, to url: URL) throws {
143 guard
144 let format = AVAudioFormat(
145 commonFormat: .pcmFormatFloat32,
146 sampleRate: Double(sampleRate),
147 channels: 2,
148 interleaved: false
149 ),
150 let buffer = AVAudioPCMBuffer(pcmFormat: format, frameCapacity: AVAudioFrameCount(stem.left.count)),
151 let channels = buffer.floatChannelData
152 else {
153 throw DemucsError.io("Could not allocate an audio buffer.")
154 }
155 buffer.frameLength = AVAudioFrameCount(stem.left.count)
156 for (index, samples) in [stem.left, stem.right].enumerated() {
157 samples.withUnsafeBufferPointer { source in
158 guard let base = source.baseAddress else { return }
159 channels[index].update(from: base, count: min(source.count, stem.left.count))
160 }
161 }
162
163 try? FileManager.default.removeItem(at: url)
164 let file = try AVAudioFile(forWriting: url, settings: format.settings)
165 try file.write(from: buffer)
166 }
167 }