adding support for multiple channels for audio ai workflows

audiosr_runner.py

@@ -1,4 +1,5 @@
 import argparse
+import json
 import os
 import subprocess
 import sys
@@ -25,12 +26,45 @@ def run_ffmpeg(args, error_prefix):
         raise RuntimeError("{}: {}".format(error_prefix, err))
 
 
+def probe_audio_info(path):
+    cmd = [
+        "ffprobe",
+        "-v",
+        "error",
+        "-select_streams",
+        "a:0",
+        "-show_entries",
+        "stream=sample_rate,channels,channel_layout",
+        "-of",
+        "json",
+        path,
+    ]
+    try:
+        result = subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+        data = json.loads(result.stdout or "{}")
+    except Exception:
+        return {}
+    streams = data.get("streams") or []
+    if not streams:
+        return {}
+    stream = streams[0] or {}
+    return {
+        "sample_rate": int(stream.get("sample_rate") or 0),
+        "channels": int(stream.get("channels") or 0),
+        "channel_layout": str(stream.get("channel_layout") or "").strip(),
+    }
+
+
 def decode_audio_to_wav(source_path, wav_path):
     run_ffmpeg(["-i", source_path, "-vn", "-acodec", "pcm_s16le", wav_path], "ffmpeg audio decode failed")
 
 
-def encode_audio_to_flac(source_path, flac_path):
-    run_ffmpeg(["-i", source_path, "-vn", "-c:a", "flac", flac_path], "ffmpeg FLAC encode failed")
+def encode_audio_to_flac(source_path, flac_path, channel_layout=None):
+    args = ["-i", source_path, "-vn"]
+    if str(channel_layout or "").strip():
+        args.extend(["-channel_layout", str(channel_layout).strip()])
+    args.extend(["-c:a", "flac", flac_path])
+    run_ffmpeg(args, "ffmpeg FLAC encode failed")
 
 
 def load_pcm16_wav(path):
@@ -61,6 +95,11 @@ def save_pcm16_wav(path, audio, sample_rate):
         wav_file.writeframes(pcm.tobytes())
 
 
+def save_mono_pcm16_wav(path, audio, sample_rate):
+    audio = np.asarray(audio, dtype=np.float32).reshape(1, -1)
+    save_pcm16_wav(path, audio, sample_rate)
+
+
 def patch_torchaudio_load():
     import torch
     import torchaudio
@@ -109,6 +148,46 @@ def run_audiosr(build_model, super_resolution, source_wav, model_name, device_na
             pass
 
 
+def normalize_audiosr_output(waveform):
+    out_np = np.asarray(waveform, dtype=np.float32)
+    if out_np.ndim == 3:
+        out_np = out_np[0]
+    if out_np.ndim == 1:
+        out_np = out_np.reshape(1, -1)
+    return out_np
+
+
+def run_audiosr_with_model(model, super_resolution, source_wav, device_name):
+    log("Running AudioSR super resolution on {}".format(device_name))
+    return super_resolution(
+        model,
+        source_wav,
+        seed=42,
+        guidance_scale=3.5,
+        ddim_steps=50,
+        latent_t_per_second=12.8,
+    )
+
+
+def run_audiosr_channels(build_model, super_resolution, source_audio_np, source_sr, model_name, device_name, tmp_dir):
+    log("Loading AudioSR model '{}' on {}".format(model_name, device_name))
+    model = build_model(model_name=model_name, device=device_name)
+    try:
+        channel_outputs = []
+        for channel_index in range(int(source_audio_np.shape[0])):
+            log("Enhancing channel {}/{} on {}".format(channel_index + 1, int(source_audio_np.shape[0]), device_name))
+            channel_path = os.path.join(tmp_dir, "channel_{}.wav".format(channel_index))
+            save_mono_pcm16_wav(channel_path, source_audio_np[channel_index], int(source_sr))
+            waveform = run_audiosr_with_model(model, super_resolution, channel_path, device_name)
+            channel_outputs.append(normalize_audiosr_output(waveform)[0])
+        return np.stack(channel_outputs, axis=0)
+    finally:
+        try:
+            model.cpu()
+        except Exception:
+            pass
+
+
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("--input", required=True)
@@ -122,6 +201,7 @@ def main():
     input_path = os.path.abspath(args.input)
     output_path = os.path.abspath(args.output)
     os.makedirs(os.path.dirname(output_path), exist_ok=True)
+    source_info = probe_audio_info(input_path)
 
     tmp_dir = tempfile.mkdtemp(prefix="openshot_audiosr_")
     try:
@@ -129,6 +209,7 @@ def main():
         enhanced_wav = os.path.join(tmp_dir, "enhanced.wav")
         log("Decoding input audio with ffmpeg")
         decode_audio_to_wav(input_path, source_wav)
+        source_audio_np, source_sr = load_pcm16_wav(source_wav)
 
         patch_torchaudio_load()
         log("Importing AudioSR")
@@ -138,7 +219,7 @@ def main():
 
         waveform = None
         try:
-            waveform = run_audiosr(build_model, super_resolution, source_wav, args.model_name, "auto")
+            waveform = run_audiosr_channels(build_model, super_resolution, source_audio_np, source_sr, args.model_name, "auto", tmp_dir)
         except Exception as ex:
             if not is_cuda_oom(ex):
                 raise
@@ -148,16 +229,12 @@ def main():
                     torch.cuda.empty_cache()
             except Exception:
                 pass
-            waveform = run_audiosr(build_model, super_resolution, source_wav, args.model_name, "cpu")
+            waveform = run_audiosr_channels(build_model, super_resolution, source_audio_np, source_sr, args.model_name, "cpu", tmp_dir)
 
-        out_np = np.asarray(waveform, dtype=np.float32)
-        if out_np.ndim == 3:
-            out_np = out_np[0]
-        if out_np.ndim == 1:
-            out_np = out_np.reshape(1, -1)
+        out_np = normalize_audiosr_output(waveform)
         log("Encoding enhanced audio to FLAC")
         save_pcm16_wav(enhanced_wav, out_np, 48000)
-        encode_audio_to_flac(enhanced_wav, output_path)
+        encode_audio_to_flac(enhanced_wav, output_path, source_info.get("channel_layout"))
         log("AudioSR output ready: {}".format(output_path))
         return 0
     finally:

deepfilternet_runner.py

@@ -1,4 +1,5 @@
 import argparse
+import json
 import os
 import subprocess
 import sys
@@ -7,6 +8,7 @@ import types
 import wave
 
 import numpy as np
+import torch
 
 
 def log(message):
@@ -60,12 +62,45 @@ def run_ffmpeg(args, error_prefix):
         raise RuntimeError("{}: {}".format(error_prefix, err))
 
 
+def probe_audio_info(path):
+    cmd = [
+        "ffprobe",
+        "-v",
+        "error",
+        "-select_streams",
+        "a:0",
+        "-show_entries",
+        "stream=sample_rate,channels,channel_layout",
+        "-of",
+        "json",
+        path,
+    ]
+    try:
+        result = subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+        data = json.loads(result.stdout or "{}")
+    except Exception:
+        return {}
+    streams = data.get("streams") or []
+    if not streams:
+        return {}
+    stream = streams[0] or {}
+    return {
+        "sample_rate": int(stream.get("sample_rate") or 0),
+        "channels": int(stream.get("channels") or 0),
+        "channel_layout": str(stream.get("channel_layout") or "").strip(),
+    }
+
+
 def decode_audio_to_wav(source_path, wav_path):
     run_ffmpeg(["-i", source_path, "-vn", "-acodec", "pcm_s16le", wav_path], "ffmpeg audio decode failed")
 
 
-def encode_audio_to_flac(source_path, flac_path):
-    run_ffmpeg(["-i", source_path, "-vn", "-c:a", "flac", flac_path], "ffmpeg FLAC encode failed")
+def encode_audio_to_flac(source_path, flac_path, channel_layout=None):
+    args = ["-i", source_path, "-vn"]
+    if str(channel_layout or "").strip():
+        args.extend(["-channel_layout", str(channel_layout).strip()])
+    args.extend(["-c:a", "flac", flac_path])
+    run_ffmpeg(args, "ffmpeg FLAC encode failed")
 
 
 def load_pcm16_wav(path):
@@ -110,6 +145,20 @@ def match_audio_length(audio, target_length):
     return F.pad(audio, (0, target_length - current))
 
 
+def enhance_channel(model, df_state, ta_functional, df_enhance, channel_audio, source_sr, amount):
+    model_sr = int(getattr(df_state, "sr", 48000)() if callable(getattr(df_state, "sr", None)) else getattr(df_state, "sr", 48000))
+    work_audio = channel_audio
+    if int(source_sr) != model_sr:
+        work_audio = ta_functional.resample(work_audio, int(source_sr), model_sr)
+    enhanced_audio = df_enhance(model, df_state, work_audio, pad=True)
+    enhanced_audio = (work_audio * (1.0 - amount)) + (enhanced_audio * amount)
+    enhanced_audio = torch.clamp(enhanced_audio, -1.0, 1.0)
+    if int(source_sr) != model_sr:
+        enhanced_audio = ta_functional.resample(enhanced_audio, model_sr, int(source_sr))
+    enhanced_audio = match_audio_length(enhanced_audio, int(channel_audio.shape[-1]))
+    return enhanced_audio
+
+
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("--input", required=True)
@@ -130,6 +179,7 @@ def main():
     output_path = os.path.abspath(args.output)
     amount = float(max(0.0, min(1.0, args.amount)))
     os.makedirs(os.path.dirname(output_path), exist_ok=True)
+    source_info = probe_audio_info(input_path)
 
     tmp_dir = tempfile.mkdtemp(prefix="openshot_df_runner_")
     try:
@@ -144,7 +194,7 @@ def main():
         if amount <= 0.0:
             log("Noise reduction is 0.0, copying input audio to FLAC")
             save_pcm16_wav(enhanced_wav, source_audio.cpu().numpy(), int(source_sr))
-            encode_audio_to_flac(enhanced_wav, output_path)
+            encode_audio_to_flac(enhanced_wav, output_path, source_info.get("channel_layout"))
             return 0
 
         log("Loading DeepFilterNet3 model")
@@ -155,23 +205,18 @@ def main():
             default_model="DeepFilterNet3",
         )
 
-        model_sr = int(getattr(df_state, "sr", 48000)() if callable(getattr(df_state, "sr", None)) else getattr(df_state, "sr", 48000))
-        work_audio = source_audio
-        if int(source_sr) != model_sr:
-            work_audio = ta_functional.resample(work_audio, int(source_sr), model_sr)
-
-        log("Running DeepFilterNet enhancement")
-        enhanced_audio = df_enhance(model, df_state, work_audio, pad=True)
-        enhanced_audio = (work_audio * (1.0 - amount)) + (enhanced_audio * amount)
-        enhanced_audio = torch.clamp(enhanced_audio, -1.0, 1.0)
-
-        if int(source_sr) != model_sr:
-            enhanced_audio = ta_functional.resample(enhanced_audio, model_sr, int(source_sr))
-        enhanced_audio = match_audio_length(enhanced_audio, int(source_audio.shape[-1]))
+        log("Running DeepFilterNet enhancement channel-by-channel")
+        enhanced_channels = []
+        for channel_index in range(int(source_audio.shape[0])):
+            log("Enhancing channel {}/{}".format(channel_index + 1, int(source_audio.shape[0])))
+            channel_audio = source_audio[channel_index : channel_index + 1, :]
+            enhanced_channel = enhance_channel(model, df_state, ta_functional, df_enhance, channel_audio, source_sr, amount)
+            enhanced_channels.append(enhanced_channel)
+        enhanced_audio = torch.cat(enhanced_channels, dim=0)
 
         log("Encoding denoised audio to FLAC")
         save_pcm16_wav(enhanced_wav, enhanced_audio.cpu().numpy(), int(source_sr))
-        encode_audio_to_flac(enhanced_wav, output_path)
+        encode_audio_to_flac(enhanced_wav, output_path, source_info.get("channel_layout"))
         log("DeepFilterNet output ready: {}".format(output_path))
 
         if args.release_model: