Update main.py

main.py

@@ -8,7 +8,6 @@ import random
 import numpy as np
 from scipy.signal.windows import hann
 import soundfile as sf
-import torch
 import librosa
 from audiosr import build_model, super_resolution
 from scipy import signal
@@ -16,25 +15,27 @@ import pyloudnorm as pyln
 import tempfile
 import spaces 
 
+
 class AudioUpscaler:
     """
     Upscales audio using the AudioSR model.
     """
 
-    def __init__(self, model_name="basic", device="auto"):
+    def __init__(self, model_name="basic", device="cuda"):
         """
         Initializes the AudioUpscaler.
 
         Args:
             model_name (str, optional): Name of the AudioSR model to use. Defaults to "basic".
-            device (str, optional): Device to use for inference. Defaults to "auto".
+            device (str, optional): Device to use for inference. Defaults to "cuda".
         """
 
         self.model_name = model_name
         self.device = device
-        self.sr = 48000
+        self.sr = 44100
         self.audiosr = None  # Model will be loaded in setup()
-
+    
+    @spaces.GPU(duration=120)
     def setup(self):
         """
         Loads the AudioSR model.
@@ -107,12 +108,12 @@ class AudioUpscaler:
         self,
         input_file,
         chunk_size=5.12,
-        overlap=0.1,
+        overlap=0.16,
         seed=None,
         guidance_scale=3.5,
         ddim_steps=50,
         multiband_ensemble=True,
-        input_cutoff=14000,
+        input_cutoff=8000,
     ):
         """
         Processes the audio in chunks and performs upsampling.
@@ -130,7 +131,7 @@ class AudioUpscaler:
         Returns:
             np.ndarray: Upsampled audio data.
         """
-
+        chunk_size = random.randint(a=0, b=10)*0.08
         audio, sr = librosa.load(input_file, sr=input_cutoff * 2, mono=False)
         audio = audio.T
         sr = input_cutoff * 2
@@ -141,12 +142,13 @@ class AudioUpscaler:
         else:
             audio_ch1 = audio
 
-        chunk_samples = int(chunk_size * sr)
+        chunk_samples   = int(chunk_size * sr)
         overlap_samples = int(overlap * chunk_samples)
 
-        output_chunk_samples = int(chunk_size * self.sr)
-        output_overlap_samples = int(overlap * output_chunk_samples)
-        enable_overlap = True if overlap > 0 else False
+
+        output_chunk_samples    = int(chunk_size * self.sr)
+        output_overlap_samples  = int(overlap * output_chunk_samples)
+        enable_overlap          = True if overlap > 0 else False
 
         def process_chunks(audio):
             chunks = []
@@ -320,7 +322,7 @@ class AudioUpscaler:
         chunk_size=10.24,
         seed=None,
         multiband_ensemble=True,
-        input_cutoff=14000,
+        input_cutoff=8000,
     ):
         """
         Upscales the audio and saves the result.
@@ -338,6 +340,7 @@ class AudioUpscaler:
         """
         if seed == 0:
             seed = random.randint(0, 2**32 - 1)
+        chunk_size = random.randint(0, 10) * 0.08
 
         os.makedirs(output_folder, exist_ok=True)
         waveform = self._process_audio(
@@ -385,7 +388,6 @@ def inference(audio_file, model_name, guidance_scale, ddim_steps, seed):
     
     return (48000, waveform)
 
-@spaces.GPU(duration=300)
 def upscale_audio(
     input_file,
     output_folder,
@@ -415,6 +417,7 @@ def upscale_audio(
         tuple: Upscaled audio data and sample rate.
     """
     torch.cuda.empty_cache()
+    chunk_size = random.randint(a=0, b=10)*0.08
 
     gc.collect()
     upscaler = AudioUpscaler()