Update Apollo/inference.py

Apollo/inference.py

@@ -8,28 +8,35 @@ import argparse
 import numpy as np
 import yaml
 from ml_collections import ConfigDict
+#from omegaconf import OmegaConf
+
 import warnings
 warnings.filterwarnings("ignore")
 
 def get_config(config_path):
     with open(config_path) as f:
+        #config = OmegaConf.load(config_path)
         config = ConfigDict(yaml.load(f, Loader=yaml.FullLoader))
         return config
 
-def load_audio(file_path, sr=44100):
-    audio, samplerate = librosa.load(file_path, mono=False, sr=sr)
+def load_audio(file_path):
+    audio, samplerate = librosa.load(file_path, mono=False, sr=44100)
     print(f'INPUT audio.shape = {audio.shape} | samplerate = {samplerate}')
+    #audio = dBgain(audio, -6)
     return torch.from_numpy(audio), samplerate
 
 def save_audio(file_path, audio, samplerate=44100):
+    #audio = dBgain(audio, +6)
     sf.write(file_path, audio.T, samplerate, subtype="PCM_16")
 
-def process_chunk(chunk, model, device):
-    chunk = chunk.unsqueeze(0).to(device)
+def process_chunk(chunk):
+    chunk = chunk.unsqueeze(0).cuda()
     with torch.no_grad():
         return model(chunk).squeeze(0).squeeze(0).cpu()
 
 def _getWindowingArray(window_size, fade_size):
+    # IMPORTANT NOTE :
+    # no fades here in the end, only removing the failed ending of the chunk
     fadein = torch.linspace(1, 1, fade_size)
     fadeout = torch.linspace(0, 0, fade_size)
     window = torch.ones(window_size)
@@ -42,26 +49,28 @@ def dBgain(audio, volume_gain_dB):
     gained_audio = audio * gain 
     return gained_audio
 
-def main(input_wav, output_wav, ckpt_path, feature_dim, layer, sr, win, chunk_size, overlap):
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    print(f"Using device: {device}")
 
-    # Modeli yükle
-    model = look2hear.models.BaseModel.from_pretrain(
-        ckpt_path, sr=sr, win=win, feature_dim=feature_dim, layer=layer
-    ).to(device)
+def main(input_wav, output_wav, ckpt_path):
+    os.environ['CUDA_VISIBLE_DEVICES'] = "0"
+
+    global model
+    feature_dim = config['model']['feature_dim']
+    sr = config['model']['sr']
+    win = config['model']['win']
+    layer = config['model']['layer']
+    model = look2hear.models.BaseModel.from_pretrain(ckpt_path, sr=sr, win=win, feature_dim=feature_dim, layer=layer).cuda()
 
-    test_data, samplerate = load_audio(input_wav, sr=sr)
+    test_data, samplerate = load_audio(input_wav)
     
     C = chunk_size * samplerate  # chunk_size seconds to samples
     N = overlap
     step = C // N
-    fade_size = 3 * samplerate  # 3 seconds
+    fade_size = 3 * 44100 # 3 seconds
     print(f"N = {N} | C = {C} | step = {step} | fade_size = {fade_size}")
     
     border = C - step
     
-    # Handle mono inputs correctly
+    # handle mono inputs correctly
     if len(test_data.shape) == 1:
         test_data = test_data.unsqueeze(0) 
 
@@ -86,7 +95,7 @@ def main(input_wav, output_wav, ckpt_path, feature_dim, layer, sr, win, chunk_si
             else:
                 part = torch.nn.functional.pad(input=part, pad=(0, C - length, 0, 0), mode='constant', value=0)
 
-        out = process_chunk(part, model, device)
+        out = process_chunk(part)
 
         window = windowingArray
         if i == 0:  # First audio chunk, no fadein
@@ -122,24 +131,20 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Audio Inference Script")
     parser.add_argument("--in_wav", type=str, required=True, help="Path to input wav file")
     parser.add_argument("--out_wav", type=str, required=True, help="Path to output wav file")
-    parser.add_argument("--ckpt", type=str, required=True, help="Path to model checkpoint file")
-    parser.add_argument("--config", type=str, required=True, help="Path to model config file")
-    parser.add_argument("--chunk_size", type=int, default=10, help="Chunk size value in seconds")
-    parser.add_argument("--overlap", type=int, default=2, help="Overlap")
-    parser.add_argument("--feature_dim", type=int, default=256, help="Feature dimension")
-    parser.add_argument("--layer", type=int, default=6, help="Number of layers")
-    parser.add_argument("--sr", type=int, default=44100, help="Sample rate")
-    parser.add_argument("--win", type=int, default=20, help="Window size")
+    parser.add_argument("--ckpt", type=str, required=True, help="Path to model checkpoint file", default="model/pytorch_model.bin")
+    parser.add_argument("--config", type=str, help="Path to model config file", default="configs/apollo.yaml")
+    parser.add_argument("--chunk_size", type=int, help="chunk size value in seconds", default=10)
+    parser.add_argument("--overlap", type=int, help="Overlap", default=2)
     args = parser.parse_args()
     
+    ckpt_path = args.ckpt
+    chunk_size = args.chunk_size
+    overlap = args.overlap
     config = get_config(args.config)
     print(config['model'])
-    print(f'ckpt_path = {args.ckpt}')
-    print(f'chunk_size = {args.chunk_size}, overlap = {args.overlap}')
-    print(f'feature_dim = {args.feature_dim}, layer = {args.layer}, sr = {args.sr}, win = {args.win}')
-
-    main(
-        args.in_wav, args.out_wav, args.ckpt, 
-        args.feature_dim, args.layer, args.sr, args.win, 
-        args.chunk_size, args.overlap
-    )
+    print(f'ckpt_path = {ckpt_path}')
+    #print(f'config = {config}')
+    print(f'chunk_size = {chunk_size}, overlap = {overlap}')
+    
+
+    main(args.in_wav, args.out_wav, ckpt_path)