Shape fix

app.py

@@ -3,6 +3,7 @@ from audioseal import AudioSeal
 import gradio as gr
 import soundfile as sf
 import librosa
+import torch
 
 from pyharp.core import ModelCard, build_endpoint
 from pyharp.media.audio import save_audio
@@ -19,24 +20,28 @@ model_card = ModelCard(
 print("Initializing AudioSeal model...")
 generator = AudioSeal.load_generator("audioseal_wm_16bits")
 detector = AudioSeal.load_detector("audioseal_detector_16bits")
+generator.eval()
+detector.eval()
 
 def load_audio(audio_path):
     try:
         wav, sr = librosa.load(audio_path, sr=16000, mono=True)
-        return wav
+        return wav[None, None, :]
     
     except Exception as e:
         print(f"Audio preprocessing failed: {e}")
         raise ValueError(f"Failed to load audio: {str(e)}")
 
 def process_fn(inp_audio, option_text):
-    wav = load_audio(inp_audio)
+    wav = torch.from_numpy(load_audio(inp_audio))
     sr = 16000
     if option_text == "Generate Watermark":
-        watermark = generator.get_watermark(wav, sr)
-        output_audio = watermark + wav
-        sig = AudioSignal(output_audio.astype("float32"), sample_rate=sr)
-        return save_audio(sig)
+        with torch.no_grad():
+            watermark = generator.get_watermark(wav, sr)
+            output_audio = watermark + wav
+            output_audio = output_audio.detach().cpu().numpy()[0]
+            sig = AudioSignal(output_audio.astype("float32"), sample_rate=sr)
+            return save_audio(sig)
     else:
         result, message = detector.detect_watermark(wav, sr)