add multiband ensemble control, clean up debug prints

app.py

@@ -1,30 +1,20 @@
 import sys
-print("Script started", flush=True)
-print(f"Python version: {sys.version}", flush=True)
 import os
-print("os imported", flush=True)
 import tempfile
-print("tempfile imported", flush=True)
 import soundfile as sf
-print("soundfile imported", flush=True)
 import torch
-print("torch imported", flush=True)
 from audiosr import build_model, super_resolution_long_audio
-print("audiosr imported", flush=True)
 from pyharp import ModelCard, build_endpoint
-print("pyharp imported", flush=True)
 import gradio as gr
-print("gradio imported", flush=True)
 import random
-print("random imported", flush=True)
+import numpy as np
+import librosa
+from scipy import signal
 
-print("Setting env vars", flush=True)
 os.environ["TOKENIZERS_PARALLELISM"] = "true"
 torch.set_float32_matmul_precision("high")
-print("Env vars set", flush=True)
 
 model = None
-print("Model variable initialized", flush=True)
 
 def get_model():
     global model
@@ -34,16 +24,38 @@ def get_model():
         print("Model loaded.", flush=True)
     return model
 
+def match_array_shapes(array_1: np.ndarray, array_2: np.ndarray):
+    if (len(array_1.shape) == 1) & (len(array_2.shape) == 1):
+        if array_1.shape[0] > array_2.shape[0]:
+            array_1 = array_1[:array_2.shape[0]]
+        elif array_1.shape[0] < array_2.shape[0]:
+            array_1 = np.pad(array_1, ((array_2.shape[0] - array_1.shape[0], 0)), 'constant', constant_values=0)
+    else:
+        if array_1.shape[1] > array_2.shape[1]:
+            array_1 = array_1[:,:array_2.shape[1]]
+        elif array_1.shape[1] < array_2.shape[1]:
+            padding = array_2.shape[1] - array_1.shape[1]
+            array_1 = np.pad(array_1, ((0,0), (0,padding)), 'constant', constant_values=0)
+    return array_1
+
+def lr_filter(audio, cutoff, filter_type, order=12, sr=48000):
+    audio = audio.T
+    nyquist = 0.5 * sr
+    normal_cutoff = cutoff / nyquist
+    b, a = signal.butter(order//2, normal_cutoff, btype=filter_type, analog=False)
+    sos = signal.tf2sos(b, a)
+    filtered_audio = signal.sosfiltfilt(sos, audio)
+    return filtered_audio.T
+
 model_card = ModelCard(
     name="AudioSR",
     description="Upsample any audio to 48kHz using audio super-resolution.",
     author="Haohe Liu, Ke Chen, Qiao Tian, Wenwu Wang, Mark D. Plumbley",
     tags=["audio", "super-resolution", "upsampling", "48kHz"],
 )
-print("Model card created", flush=True)
 
 @torch.inference_mode()
-def process_fn(input_audio_path: str, ddim_steps: int, guidance_scale: float, seed: str) -> str:
+def process_fn(input_audio_path: str, ddim_steps: int, guidance_scale: float, seed: str, multiband_ensemble: bool, input_cutoff: int) -> str:
     try:
         seed_val = int(seed) if seed and seed.strip() not in ("0", "") else random.randint(1, 2**32 - 1)
     except (TypeError, ValueError):
@@ -59,17 +71,26 @@ def process_fn(input_audio_path: str, ddim_steps: int, guidance_scale: float, se
 
     output = waveform.cpu().numpy()
 
-    if output.shape[0] == 1:
-        output = output.squeeze(0)
+    if multiband_ensemble:
+        crossover_freq = int(input_cutoff) - 1000
+        low, _ = librosa.load(input_audio_path, sr=48000, mono=True)
+        out = output.squeeze() if output.ndim > 1 else output
+        out = match_array_shapes(out, low)
+        low = lr_filter(low, crossover_freq, 'lowpass', order=10)
+        high = lr_filter(out, crossover_freq, 'highpass', order=10)
+        high = lr_filter(high, 23000, 'lowpass', order=2)
+        output = low + high
     else:
-        output = output.T
+        if output.shape[0] == 1:
+            output = output.squeeze(0)
+        else:
+            output = output.T
 
     output_path = tempfile.mktemp(suffix=".wav")
     sf.write(output_path, output, samplerate=48000)
     return output_path
 
 with gr.Blocks() as demo:
-    print("Gradio blocks created", flush=True)
     input_components = [
         gr.Audio(type="filepath", label="Input Audio").harp_required(True),
         gr.Slider(minimum=10, maximum=500, step=10, value=50,
@@ -79,6 +100,11 @@ with gr.Blocks() as demo:
                   label="Guidance Scale",
                   info="Higher values follow the conditioning more closely"),
         gr.Textbox(value="0", label="Seed", info="0 = random seed"),
+        gr.Checkbox(value=False, label="Multiband Ensemble",
+                    info="Blend original low frequencies with upsampled highs"),
+        gr.Slider(minimum=4000, maximum=20000, step=1000, value=12000,
+                  label="Input Cutoff (Hz)",
+                  info="Crossover frequency for multiband ensemble"),
     ]
     output_components = [
         gr.Audio(type="filepath", label="Output Audio (48kHz)").set_info("Audio upsampled to 48kHz."),
@@ -89,7 +115,6 @@ with gr.Blocks() as demo:
         output_components=output_components,
         process_fn=process_fn,
     )
-print("Endpoint built", flush=True)
 
 print("Launching Gradio...", flush=True)
-demo.queue().launch(server_name="0.0.0.0", server_port=7860, show_error=True, pwa=True)
+demo.queue().launch(server_name="0.0.0.0", server_port=7860, show_error=True, pwa=True)