diff --git a/backend/backend.proto b/backend/backend.proto index 2935d1a62..bd59159b2 100644 --- a/backend/backend.proto +++ b/backend/backend.proto @@ -1122,11 +1122,28 @@ message AudioTransformRequest { string ModelIdentity = 5; } +// One named output of a transform that produces several from a single run. +// Source separation is the case that needs it: htdemucs yields drums, bass, +// other and vocals from one pass over the input. +message AudioTransformStem { + string name = 1; // the model's own stem id, e.g. "vocals" + string dst = 2; // path of the file written for that stem +} + message AudioTransformResult { string dst = 1; int32 sample_rate = 2; int32 samples = 3; bool reference_provided = 4; + // Every named output the run produced, in the model's own order, including + // the one copied into dst. Empty for a transform with a single output. + // + // It exists because dst carries one file while separation produces several, + // and running the model once per stem would cost four full separations of + // the same audio. The backend runs once, writes each stem beside dst, and + // names them here; without this field the other stems are on disk but no + // caller can find them, which is the same as not having produced them. + repeated AudioTransformStem stems = 5; } // Bidirectional streaming audio transform. The first message MUST carry a diff --git a/backend/cpp/audio-cpp/grpc-server.cpp b/backend/cpp/audio-cpp/grpc-server.cpp index 63db70f42..57280cad3 100644 --- a/backend/cpp/audio-cpp/grpc-server.cpp +++ b/backend/cpp/audio-cpp/grpc-server.cpp @@ -896,9 +896,18 @@ public: } for (size_t i = 0; i < result.named_audio_outputs.size(); ++i) { + const std::string sibling = + audiocpp_backend::sibling_stem_path(request->dst(), names[i]); audiocpp_backend::write_audio_file( - audiocpp_backend::sibling_stem_path(request->dst(), names[i]), - result.named_audio_outputs[i].audio); + sibling, result.named_audio_outputs[i].audio); + // Named in the response, in the model's own order. Without + // this the siblings are files nobody can find, and a caller + // that wants the drums as well as the vocals has to run the + // whole separation again per stem, which is the cost the + // single run exists to avoid. + auto *stem = response->add_stems(); + stem->set_name(names[i]); + stem->set_dst(sibling); } chosen = &result.named_audio_outputs[static_cast(choice.index)] .audio; diff --git a/core/backend/audio_transform.go b/core/backend/audio_transform.go index bb006fc10..99c46a4ff 100644 --- a/core/backend/audio_transform.go +++ b/core/backend/audio_transform.go @@ -33,6 +33,22 @@ type AudioTransformOutputs struct { Dst string AudioPath string ReferencePath string + // Stems are the other named outputs the same run produced, in the model's + // own order and including the one whose content Dst carries. Empty for a + // single-output transform. + // + // A separation backend writes every stem beside Dst from ONE inference. + // Dropping them here would mean a caller who wants drums as well as vocals + // has to run the whole separation again per stem, which is precisely what + // the single run exists to avoid. + Stems []AudioTransformStem +} + +// AudioTransformStem is one named output of a multi-output transform, e.g. the +// "vocals" track of a source separation. +type AudioTransformStem struct { + Name string + Dst string } // ModelAudioTransform runs the unary AudioTransform RPC and returns the @@ -128,9 +144,40 @@ func ModelAudioTransform( Dst: dst, AudioPath: persistedAudio, ReferencePath: persistedRef, + Stems: collectStems(res, audioDir), }, res, nil } +// collectStems turns the backend's reported stems into the caller-facing list. +// +// Every path is checked to be a direct child of audioDir, the generated-content +// directory this request handed the backend. A backend is a separate process +// and its response is not this process's data: a stem path pointing at /etc or +// at another user's file would otherwise be served straight back through the +// HTTP layer, which resolves these into URLs. A stem that fails the check is +// dropped rather than fatal, so a well-behaved majority still reaches the +// caller. +func collectStems(res *proto.AudioTransformResult, audioDir string) []AudioTransformStem { + if res == nil || len(res.GetStems()) == 0 { + return nil + } + stems := make([]AudioTransformStem, 0, len(res.GetStems())) + for _, stem := range res.GetStems() { + name, path := stem.GetName(), stem.GetDst() + if name == "" || path == "" { + continue + } + if filepath.Dir(filepath.Clean(path)) != filepath.Clean(audioDir) { + continue + } + stems = append(stems, AudioTransformStem{Name: name, Dst: path}) + } + if len(stems) == 0 { + return nil + } + return stems +} + // ModelAudioTransformStream opens the bidirectional AudioTransformStream RPC // and returns the underlying stream client. The caller is responsible for // sending the initial Config message, subsequent Frame messages, and for diff --git a/core/config/backend_capabilities.go b/core/config/backend_capabilities.go index 3de119bbb..7b328cee8 100644 --- a/core/config/backend_capabilities.go +++ b/core/config/backend_capabilities.go @@ -223,6 +223,20 @@ type BackendCapability struct { AcceptsVideos bool // AcceptsAudios indicates multimodal audio input in Predict. AcceptsAudios bool + // AudioTransformInputMono16k declares that this backend's AudioTransform + // input must be folded to 16 kHz mono 16-bit WAV before it is handed over. + // + // Opt-IN, and the default of false means "hand the backend the upload as + // it is". The /audio/transform endpoint used to fold EVERY upload to + // 16 kHz mono, which is what LocalVQE wants for acoustic echo cancellation + // and what no source separation model can survive: htdemucs and + // mel_band_roformer refuse any rate but their checkpoint's own (44.1 kHz + // for every published one) and work in stereo, so every separation request + // made through the HTTP API failed with an INTERNAL raised inside the + // engine, while a direct gRPC call worked. Declaring the need rather than + // defaulting to it means a backend that wants the fold says so and a + // backend that does not needs no entry here at all. + AudioTransformInputMono16k bool // VoiceCloning describes the backend's per-request reference-audio // contract. Model variants that share a backend may narrow this further; // use VoiceCloningForModel for UI/API decisions. @@ -578,7 +592,11 @@ var BackendCapabilities = map[string]BackendCapability{ GRPCMethods: []GRPCMethod{MethodAudioTransform}, PossibleUsecases: []string{UsecaseAudioTransform}, DefaultUsecases: []string{UsecaseAudioTransform}, - Description: "LocalVQE — joint AEC, noise suppression, and dereverberation for 16 kHz mono speech", + // The model is trained on 16 kHz mono speech and its AEC needs the + // input and the loopback reference in the same shape, so the endpoint + // keeps folding uploads for this backend. + AudioTransformInputMono16k: true, + Description: "LocalVQE — joint AEC, noise suppression, and dereverberation for 16 kHz mono speech", }, // --- Utility backends --- @@ -740,6 +758,18 @@ func GetBackendCapability(backend string) *BackendCapability { return nil } +// AudioTransformRequiresMono16kInput reports whether /audio/transform must fold +// uploads to 16 kHz mono before handing them to this backend. +// +// False for an unknown backend, which is the safe answer: an unregistered +// backend gets its upload unchanged, so a model that needs the file intact +// (source separation, voice conversion at 44.1 kHz) works without an entry +// here, and one that needs the fold cannot get it by accident. +func AudioTransformRequiresMono16kInput(backend string) bool { + capability := GetBackendCapability(backend) + return capability != nil && capability.AudioTransformInputMono16k +} + // VoiceCloningForModel returns the reference-audio contract only when the // installed model variant can honor it. Several backends serve both Base // (voice cloning) and CustomVoice/VoiceDesign models, so backend name alone is diff --git a/core/config/backend_capabilities_test.go b/core/config/backend_capabilities_test.go index 1a5e1d7db..8f22ad44a 100644 --- a/core/config/backend_capabilities_test.go +++ b/core/config/backend_capabilities_test.go @@ -72,6 +72,33 @@ var _ = Describe("GetBackendCapability", func() { }) }) +// The fold to 16 kHz mono in /audio/transform is opt-IN. It used to be +// unconditional, which made source separation unreachable through the HTTP +// API: htdemucs and mel_band_roformer refuse any rate but their checkpoint's +// own and separate using the stereo image, so every such request died with an +// INTERNAL raised inside the engine while the same call over gRPC worked. +var _ = Describe("AudioTransformRequiresMono16kInput", func() { + It("folds for localvqe, whose AEC is trained on 16 kHz mono", func() { + Expect(AudioTransformRequiresMono16kInput("localvqe")).To(BeTrue()) + }) + + It("does not fold for an unregistered backend", func() { + Expect(AudioTransformRequiresMono16kInput("audio-cpp")).To(BeFalse()) + Expect(AudioTransformRequiresMono16kInput("nonexistent")).To(BeFalse()) + Expect(AudioTransformRequiresMono16kInput("")).To(BeFalse()) + }) + + It("is claimed by no other registered backend", func() { + for name, capability := range BackendCapabilities { + if name == "localvqe" { + continue + } + Expect(capability.AudioTransformInputMono16k).To(BeFalse(), + "backend %q asks for the 16 kHz mono fold; that has to be a deliberate, documented need", name) + } + }) +}) + var _ = Describe("VoiceCloningForModel", func() { voiceCloningSetting := func(enabled bool) *bool { return &enabled } diff --git a/core/http/endpoints/localai/audio_transform.go b/core/http/endpoints/localai/audio_transform.go index 2868afed6..8123999e6 100644 --- a/core/http/endpoints/localai/audio_transform.go +++ b/core/http/endpoints/localai/audio_transform.go @@ -105,14 +105,24 @@ func AudioTransformEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, } defer func() { _ = os.RemoveAll(dir) }() - audioPath, err := saveMultipartFileAsWAV(audioFile, dir, "audio") + // Whether the upload is folded to 16 kHz mono is the BACKEND'S + // declaration and not this endpoint's default. LocalVQE's echo + // cancellation needs that shape; source separation cannot survive it. + // See config.AudioTransformRequiresMono16kInput. + fold := config.AudioTransformRequiresMono16kInput(cfg.Backend) + + audioPath, err := saveMultipartFileAsWAV(audioFile, dir, "audio", fold) if err != nil { return err } var referencePath string if refFile, err := c.FormFile("reference"); err == nil { - referencePath, err = saveMultipartFileAsWAV(refFile, dir, "reference") + // The reference is folded on the same terms as the primary input. + // For AEC the two have to be the same shape to line up sample for + // sample; for voice conversion the reference is a speaker clip + // whose rate the family resolves for itself. + referencePath, err = saveMultipartFileAsWAV(refFile, dir, "reference", fold) if err != nil { return err } @@ -154,7 +164,7 @@ func AudioTransformEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, // history alongside the output. The /generated-audio/ prefix is // the same one ttsApi uses (parsed from Content-Disposition). if name := filepath.Base(out.AudioPath); name != "" { - c.Response().Header().Set(echo.HeaderAccessControlExposeHeaders, "X-Audio-Input-Url, X-Audio-Reference-Url") + c.Response().Header().Set(echo.HeaderAccessControlExposeHeaders, exposedAudioTransformHeaders) c.Response().Header().Set("X-Audio-Input-Url", "/generated-audio/"+name) } if out.ReferencePath != "" { @@ -162,6 +172,15 @@ func AudioTransformEndpoint(cl *config.ModelConfigLoader, ml *model.ModelLoader, c.Response().Header().Set("X-Audio-Reference-Url", "/generated-audio/"+name) } } + // The body can carry one file, and a separation produces several from + // one run. The rest are named here, as JSON, so a caller can fetch the + // stems it did not ask for without paying for another separation. + // Header rather than body for the same reason the input URLs are + // headers: the body is the audio itself. + if header := stemsHeader(out.Stems); header != "" { + c.Response().Header().Set(echo.HeaderAccessControlExposeHeaders, exposedAudioTransformHeaders) + c.Response().Header().Set("X-Audio-Stems", header) + } return c.Attachment(dst, filepath.Base(dst)) } } @@ -318,12 +337,60 @@ func buildConfigRequest(fmt_ proto.AudioTransformStreamConfig_SampleFormat, ctrl } } -// saveMultipartFileAsWAV materialises an uploaded multipart file into `dir` -// and converts it to LocalVQE's required shape (16 kHz mono s16 WAV) via -// ffmpeg. The conversion is a passthrough when the upload already matches. -// `name` is used as the base filename for the converted output so the dir -// stays readable for debugging (e.g. "audio.wav", "reference.wav"). -func saveMultipartFileAsWAV(fh *multipart.FileHeader, dir, name string) (string, error) { +// exposedAudioTransformHeaders is the CORS allow-list for the artifact headers +// this endpoint sets. A browser cannot read any of them without it. +const exposedAudioTransformHeaders = "X-Audio-Input-Url, X-Audio-Reference-Url, X-Audio-Stems" + +// stemsHeader renders the extra named outputs as a compact JSON array for the +// X-Audio-Stems response header: +// +// [{"name":"vocals","url":"/generated-audio/transform-1.vocals.wav"}, ...] +// +// JSON rather than a name=url list because a stem name is the MODEL'S string: +// it comes out of the checkpoint's config, so a name containing a comma or an +// equals sign would silently corrupt any hand-rolled separator format. Header +// values must also stay on one line, which json.Marshal guarantees since it +// escapes every control character. +func stemsHeader(stems []backend.AudioTransformStem) string { + if len(stems) == 0 { + return "" + } + type stemEntry struct { + Name string `json:"name"` + URL string `json:"url"` + } + entries := make([]stemEntry, 0, len(stems)) + for _, stem := range stems { + name := filepath.Base(stem.Dst) + if name == "" || name == "." || name == string(filepath.Separator) { + continue + } + entries = append(entries, stemEntry{Name: stem.Name, URL: "/generated-audio/" + name}) + } + if len(entries) == 0 { + return "" + } + encoded, err := json.Marshal(entries) + if err != nil { + // Unreachable for a slice of plain strings, and a failure here must + // not cost the caller the audio it did ask for. + xlog.Debug("audio_transform: cannot encode stem header", "error", err) + return "" + } + return string(encoded) +} + +// saveMultipartFileAsWAV materialises an uploaded multipart file into `dir` as +// a 16-bit PCM WAV. `name` is used as the base filename for the converted +// output so the dir stays readable for debugging (e.g. "audio.wav", +// "reference.wav"). +// +// `foldToMono16k` picks the target shape. True folds to 16 kHz mono, which is +// what LocalVQE requires; false keeps the upload's own rate and channel count, +// which is what everything else needs and what source separation cannot work +// without. Either way a WAV that already matches is passed through rather than +// re-encoded. +func saveMultipartFileAsWAV(fh *multipart.FileHeader, dir, name string, foldToMono16k bool) (string, error) { f, err := fh.Open() if err != nil { return "", err @@ -342,7 +409,11 @@ func saveMultipartFileAsWAV(fh *multipart.FileHeader, dir, name string) (string, _ = out.Close() dst := filepath.Join(dir, name+".wav") - if err := utils.AudioToWav(raw, dst); err != nil { + convert := utils.AudioToWavPreservingShape + if foldToMono16k { + convert = utils.AudioToWav + } + if err := convert(raw, dst); err != nil { return "", fmt.Errorf("normalize %s: %w", name, err) } return dst, nil diff --git a/pkg/utils/ffmpeg.go b/pkg/utils/ffmpeg.go index 1ebcc11a8..06164a663 100644 --- a/pkg/utils/ffmpeg.go +++ b/pkg/utils/ffmpeg.go @@ -34,6 +34,27 @@ func AudioToWav(src, dst string) error { return convertWithFFmpeg(src, dst) } +// AudioToWavPreservingShape converts audio to 16-bit PCM WAV while KEEPING the +// source sample rate and channel count. A WAV that is already 16-bit PCM is +// passed through byte for byte, whatever its rate or channel count. +// +// It is the counterpart to AudioToWav, which folds everything to 16 kHz mono. +// That fold is right for speech backends and destructive for the rest: source +// separation models refuse any rate but their checkpoint's own, and separate a +// centred vocal from a wide mix using the stereo image, so a 16 kHz mono +// downmix removes both the format they accept and the cue they work from. +// Callers pick between the two from the BACKEND'S declared need +// (config.AudioTransformRequiresMono16kInput), never from the file. +// +// Non-WAV uploads are still transcoded, since backends read WAV. Only the rate +// and the channel layout survive that. +func AudioToWavPreservingShape(src, dst string) error { + if strings.HasSuffix(src, ".wav") && isPCM16Wav(src) { + return passthroughWAV(src, dst) + } + return convertWithFFmpegPreservingShape(src, dst) +} + func passthroughWAV(src, dst string) error { if err := os.Link(src, dst); err == nil { return nil @@ -82,6 +103,35 @@ func convertWithFFmpeg(src, dst string) error { return nil } +// isPCM16Wav returns true when src is a valid 16-bit PCM WAV at ANY sample rate +// and ANY channel count. Deliberately weaker than isTargetWav: it is the "no +// conversion needed" test for callers that want the file's own shape kept. +func isPCM16Wav(src string) bool { + f, err := os.Open(src) + if err != nil { + return false + } + defer func() { _ = f.Close() }() + + dec := wav.NewDecoder(f) + if !dec.IsValidFile() { + return false + } + return dec.BitDepth == 16 +} + +// convertWithFFmpegPreservingShape transcodes to 16-bit PCM WAV with no -ar and +// no -ac, so ffmpeg keeps the source rate and channel count. Dropping those two +// flags is the whole difference from convertWithFFmpeg. +func convertWithFFmpegPreservingShape(src, dst string) error { + commandArgs := []string{"-y", "-i", src, "-acodec", "pcm_s16le", dst} + out, err := ffmpegCommand(commandArgs) + if err != nil { + return fmt.Errorf("error: %w out: %s", err, out) + } + return nil +} + // AudioResample resamples an audio file to the given sample rate using ffmpeg. // If sampleRate <= 0, it is a no-op and returns src unchanged. func AudioResample(src string, sampleRate int) (string, error) { diff --git a/pkg/utils/ffmpeg_shape_test.go b/pkg/utils/ffmpeg_shape_test.go new file mode 100644 index 000000000..176e6d5d2 --- /dev/null +++ b/pkg/utils/ffmpeg_shape_test.go @@ -0,0 +1,132 @@ +package utils_test + +import ( + "encoding/binary" + "os" + "path/filepath" + + laudio "github.com/mudler/LocalAI/pkg/audio" + . "github.com/mudler/LocalAI/pkg/utils" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +// writeWAV writes a 16-bit PCM WAV with the given shape. The samples are a +// ramp rather than silence, so a conversion that dropped its input would be +// visible as a shorter file and not just as a different header. +func writeWAV(path string, sampleRate uint32, channels uint16, frames int) { + f, err := os.Create(path) + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = f.Close() }() + + const bitsPerSample = uint16(16) + blockAlign := channels * (bitsPerSample / 8) + total := frames * int(channels) + dataSize := uint32(total) * uint32(bitsPerSample/8) + + hdr := laudio.WAVHeader{ + ChunkID: [4]byte{'R', 'I', 'F', 'F'}, + ChunkSize: 36 + dataSize, + Format: [4]byte{'W', 'A', 'V', 'E'}, + Subchunk1ID: [4]byte{'f', 'm', 't', ' '}, + Subchunk1Size: 16, + AudioFormat: 1, + NumChannels: channels, + SampleRate: sampleRate, + ByteRate: sampleRate * uint32(blockAlign), + BlockAlign: blockAlign, + BitsPerSample: bitsPerSample, + Subchunk2ID: [4]byte{'d', 'a', 't', 'a'}, + Subchunk2Size: dataSize, + } + Expect(binary.Write(f, binary.LittleEndian, &hdr)).To(Succeed()) + for i := 0; i < total; i++ { + Expect(binary.Write(f, binary.LittleEndian, int16(200*(i%100)))).To(Succeed()) + } +} + +// readShape reports the (rate, channels, bit depth) a WAV declares. +func readShape(path string) (uint32, uint16, uint16) { + f, err := os.Open(path) + Expect(err).ToNot(HaveOccurred()) + defer func() { _ = f.Close() }() + var hdr laudio.WAVHeader + Expect(binary.Read(f, binary.LittleEndian, &hdr)).To(Succeed()) + return hdr.SampleRate, hdr.NumChannels, hdr.BitsPerSample +} + +// The distinction these specs defend: /audio/transform used to fold EVERY +// upload to 16 kHz mono, which made source separation impossible to reach +// through the HTTP API (htdemucs refuses any rate but its checkpoint's own and +// separates using the stereo image). AudioToWav keeps that fold for the +// backends that want it; AudioToWavPreservingShape is what everything else +// gets. +var _ = Describe("utils/ffmpeg shape-preserving conversion", func() { + var dir string + + BeforeEach(func() { + dir = GinkgoT().TempDir() + }) + + It("passes a 44.1 kHz stereo WAV through byte for byte", func() { + src := filepath.Join(dir, "input.wav") + dst := filepath.Join(dir, "output.wav") + writeWAV(src, 44100, 2, 4410) + + Expect(AudioToWavPreservingShape(src, dst)).To(Succeed()) + + rate, channels, depth := readShape(dst) + Expect(rate).To(Equal(uint32(44100)), "the source rate must survive") + Expect(channels).To(Equal(uint16(2)), "the stereo image must survive") + Expect(depth).To(Equal(uint16(16))) + + before, err := os.ReadFile(src) + Expect(err).ToNot(HaveOccurred()) + after, err := os.ReadFile(dst) + Expect(err).ToNot(HaveOccurred()) + Expect(after).To(Equal(before), "a 16-bit PCM WAV is not re-encoded at all") + }) + + It("keeps an unusual rate and channel count that AudioToWav would destroy", func() { + src := filepath.Join(dir, "odd.wav") + preserved := filepath.Join(dir, "preserved.wav") + folded := filepath.Join(dir, "folded.wav") + writeWAV(src, 48000, 2, 4800) + + Expect(AudioToWavPreservingShape(src, preserved)).To(Succeed()) + rate, channels, _ := readShape(preserved) + Expect(rate).To(Equal(uint32(48000))) + Expect(channels).To(Equal(uint16(2))) + + // The contrast, so this spec fails if the two functions ever converge. + // AudioToWav shells out to ffmpeg here, which the unit suite cannot + // assume is installed, so the fold is only asserted when it succeeds. + if err := AudioToWav(src, folded); err == nil { + foldedRate, foldedChannels, _ := readShape(folded) + Expect(foldedRate).To(Equal(uint32(16000)), "AudioToWav still folds") + Expect(foldedChannels).To(Equal(uint16(1)), "AudioToWav still downmixes") + } + }) + + It("still passes a 16 kHz mono WAV through", func() { + src := filepath.Join(dir, "speech.wav") + dst := filepath.Join(dir, "speech-out.wav") + writeWAV(src, 16000, 1, 1600) + + Expect(AudioToWavPreservingShape(src, dst)).To(Succeed()) + + rate, channels, _ := readShape(dst) + Expect(rate).To(Equal(uint32(16000))) + Expect(channels).To(Equal(uint16(1))) + }) + + It("leaves the source file in place", func() { + src := filepath.Join(dir, "keep.wav") + dst := filepath.Join(dir, "keep-out.wav") + writeWAV(src, 44100, 2, 441) + + Expect(AudioToWavPreservingShape(src, dst)).To(Succeed()) + Expect(src).To(BeAnExistingFile(), "the upload must survive the conversion") + }) +})