Update model.py

model.py

@@ -57,69 +57,54 @@ def load_model():
 
 @torch.no_grad()
 def invert_audio(
-        model, processor, input_audio_path, out_path,
+        model, processor, input_audio, sampling_rate,
         normalize=True, flip_input=True, flip_output=False):
 
     model.config.normalize = normalize
 
-    audio_sample_1, sampling_rate_1 = audio_read(input_audio_path)
-    if sampling_rate_1 != MODEL_SAMPLING_RATE:
-        audio_sample_1 = julius.resample_frac(audio_sample_1, sampling_rate_1, MODEL_SAMPLING_RATE)
+    # Check and resample the input audio if necessary
+    if sampling_rate != MODEL_SAMPLING_RATE:
+        input_audio = julius.resample_frac(input_audio, sampling_rate, MODEL_SAMPLING_RATE)
 
-    # audio_sample [2, 9399305]
+    # Flip the audio if required
     if flip_input:
-        audio_sample_1 = torch.flip(audio_sample_1, dims=(1,))
-    
-    # pre-process the inputs
-    inputs_1 = processor(raw_audio=audio_sample_1, sampling_rate=MODEL_SAMPLING_RATE, return_tensors="pt")
+        input_audio = torch.flip(input_audio, dims=(1,))
+
+    # Pre-process the inputs
+    inputs_1 = processor(raw_audio=input_audio, sampling_rate=MODEL_SAMPLING_RATE, return_tensors="pt")
     inputs_1["input_values"] = inputs_1["input_values"].to("cuda:0")
     inputs_1["padding_mask"] = inputs_1["padding_mask"].to("cuda:0")
 
-    # explicitly encode then decode the audio inputs
+    # Explicitly encode then decode the audio inputs
     print("Encoding...")
     encoder_outputs_1 = model.encode(
         inputs_1["input_values"],
         inputs_1["padding_mask"],
         bandwidth=max(model.config.target_bandwidths))
 
-    # EMBEDDINGS (no quantized):
-    # encoder_outputs.audio_codes.shape 
-    # [216, 1, 128, 150]
-
     avg = torch.mean(encoder_outputs_1.audio_codes, (0, 3), True)
-    # [1, 1, 128, 1]
     avg_repeat = avg.repeat(
         encoder_outputs_1.audio_codes.shape[0],
         encoder_outputs_1.audio_codes.shape[1],
         1,
         encoder_outputs_1.audio_codes.shape[3])
-    # [216, 1, 128, 150]
     diff_repeat = encoder_outputs_1.audio_codes - avg_repeat
-    
-    # TODO: power factor calculations kinda useless if we keep the factor one???
+
     POWER_FACTOR = 1
     max_abs_diff = torch.max(torch.abs(diff_repeat))
     diff_abs_power = ((torch.abs(diff_repeat) / max_abs_diff) ** POWER_FACTOR) * max_abs_diff
     latents = (diff_repeat >= 0) * diff_abs_power - (diff_repeat < 0) * diff_abs_power
 
-    # difference inversion done here!
+    # Inversion of difference
     latents = latents * -1.0
 
     print("Decoding...")
     audio_values = model.decode(latents, encoder_outputs_1.audio_scales, inputs_1["padding_mask"])[0]
 
-    # [1, 2, 10264800]
     if flip_output:
         audio_values = torch.flip(audio_values, dims=(2,))
 
-    output_dir = "/home/romainpaulusisep_gmail_com/data/outputs"
+    # Return the decoded audio tensor (or NumPy array, based on your audio_write function)
     decoded_wav = audio_values.squeeze(0).to("cpu")
 
-    print("Saving output file...")
-    out_path_ = audio_write(
-        out_path,
-        sample_rate=MODEL_SAMPLING_RATE,
-        wav=decoded_wav,
-        normalize=False)
-
-    return out_path_
+    return decoded_wav