Update app.py

app.py

@@ -72,6 +72,25 @@ def compare_pitch(original_audio, processed_audio, sr=16000):
     else:
         print("Could not extract pitch from one of the signals.")
 
+def adjust_spectrogram_mean(spec, target_mean=-5.0):
+    # Calculate the current mean of the spectrogram
+    current_mean = spec.mean().item()
+    
+    # If the current mean is below the target mean, shift the values up
+    if current_mean < target_mean:
+        shift_value = target_mean - current_mean
+        print(f"Current mean: {current_mean}. Shifting by: {shift_value}")
+        
+        # Shift the entire spectrogram by the calculated shift value
+        adjusted_spec = spec + shift_value
+        
+        # Ensure that the adjusted values are still valid (in the expected range)
+        adjusted_spec = torch.clamp(adjusted_spec, min=0.0)  # Optional: prevent negative values if needed
+        return adjusted_spec
+    else:
+        print(f"Current mean: {current_mean}. No adjustment needed.")
+        return spec
+
 def infer(prompt, progress=gr.Progress(track_tqdm=True)):
     pipeline = AuffusionPipeline.from_pretrained("auffusion/auffusion")
     prompt = prompt
@@ -120,6 +139,8 @@ def infer_img2img(prompt, audio_path, desired_strength, progress=gr.Progress(tra
     else:
         print(f"✅ Spectrogram looks normal (Mean: {spec_mean_before}). No boost needed.")
     
+    spec = adjust_spectrogram_mean(spec, target_mean=-5.0)
+    
     # Normalize the spectrogram
     norm_spec = normalize_spectrogram(spec)