Format code (#366)

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

Author: github-actions[bot]

MDXNet.py

@@ -1,5 +1,5 @@
 import soundfile as sf
-import torch,pdb,time,argparse,os,warnings,sys,librosa
+import torch, pdb, time, argparse, os, warnings, sys, librosa
 import numpy as np
 import onnxruntime as ort
 from scipy.io.wavfile import write
@@ -8,96 +8,133 @@
 import torch.nn as nn
 
 dim_c = 4
-class Conv_TDF_net_trim():
-    def __init__(self, device, model_name, target_name,
-                 L, dim_f, dim_t, n_fft, hop=1024):
+
+
+class Conv_TDF_net_trim:
+    def __init__(
+        self, device, model_name, target_name, L, dim_f, dim_t, n_fft, hop=1024
+    ):
         super(Conv_TDF_net_trim, self).__init__()
 
         self.dim_f = dim_f
-        self.dim_t = 2 ** dim_t
+        self.dim_t = 2**dim_t
         self.n_fft = n_fft
         self.hop = hop
         self.n_bins = self.n_fft // 2 + 1
         self.chunk_size = hop * (self.dim_t - 1)
-        self.window = torch.hann_window(window_length=self.n_fft, periodic=True).to(device)
+        self.window = torch.hann_window(window_length=self.n_fft, periodic=True).to(
+            device
+        )
         self.target_name = target_name
-        self.blender = 'blender' in model_name
+        self.blender = "blender" in model_name
 
-        out_c = dim_c * 4 if target_name == '*' else dim_c
-        self.freq_pad = torch.zeros([1, out_c, self.n_bins - self.dim_f, self.dim_t]).to(device)
+        out_c = dim_c * 4 if target_name == "*" else dim_c
+        self.freq_pad = torch.zeros(
+            [1, out_c, self.n_bins - self.dim_f, self.dim_t]
+        ).to(device)
 
         self.n = L // 2
 
     def stft(self, x):
         x = x.reshape([-1, self.chunk_size])
-        x = torch.stft(x, n_fft=self.n_fft, hop_length=self.hop, window=self.window, center=True, return_complex=True)
+        x = torch.stft(
+            x,
+            n_fft=self.n_fft,
+            hop_length=self.hop,
+            window=self.window,
+            center=True,
+            return_complex=True,
+        )
         x = torch.view_as_real(x)
         x = x.permute([0, 3, 1, 2])
-        x = x.reshape([-1, 2, 2, self.n_bins, self.dim_t]).reshape([-1, dim_c, self.n_bins, self.dim_t])
-        return x[:, :, :self.dim_f]
+        x = x.reshape([-1, 2, 2, self.n_bins, self.dim_t]).reshape(
+            [-1, dim_c, self.n_bins, self.dim_t]
+        )
+        return x[:, :, : self.dim_f]
 
     def istft(self, x, freq_pad=None):
-        freq_pad = self.freq_pad.repeat([x.shape[0], 1, 1, 1]) if freq_pad is None else freq_pad
+        freq_pad = (
+            self.freq_pad.repeat([x.shape[0], 1, 1, 1])
+            if freq_pad is None
+            else freq_pad
+        )
         x = torch.cat([x, freq_pad], -2)
-        c = 4 * 2 if self.target_name == '*' else 2
-        x = x.reshape([-1, c, 2, self.n_bins, self.dim_t]).reshape([-1, 2, self.n_bins, self.dim_t])
+        c = 4 * 2 if self.target_name == "*" else 2
+        x = x.reshape([-1, c, 2, self.n_bins, self.dim_t]).reshape(
+            [-1, 2, self.n_bins, self.dim_t]
+        )
         x = x.permute([0, 2, 3, 1])
         x = x.contiguous()
         x = torch.view_as_complex(x)
-        x = torch.istft(x, n_fft=self.n_fft, hop_length=self.hop, window=self.window, center=True)
+        x = torch.istft(
+            x, n_fft=self.n_fft, hop_length=self.hop, window=self.window, center=True
+        )
         return x.reshape([-1, c, self.chunk_size])
+
+
 def get_models(device, dim_f, dim_t, n_fft):
     return Conv_TDF_net_trim(
         device=device,
-        model_name='Conv-TDF', target_name='vocals',
+        model_name="Conv-TDF",
+        target_name="vocals",
         L=11,
-        dim_f=dim_f, dim_t=dim_t,
-        n_fft=n_fft
+        dim_f=dim_f,
+        dim_t=dim_t,
+        n_fft=n_fft,
     )
 
+
 warnings.filterwarnings("ignore")
-cpu = torch.device('cpu')
-device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
+cpu = torch.device("cpu")
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
 
 class Predictor:
-    def __init__(self,args):
-        self.args=args
-        self.model_ = get_models(device=cpu, dim_f=args.dim_f, dim_t=args.dim_t, n_fft=args.n_fft)
-        self.model = ort.InferenceSession(os.path.join(args.onnx,self.model_.target_name+'.onnx'), providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
-        print('onnx load done')
+    def __init__(self, args):
+        self.args = args
+        self.model_ = get_models(
+            device=cpu, dim_f=args.dim_f, dim_t=args.dim_t, n_fft=args.n_fft
+        )
+        self.model = ort.InferenceSession(
+            os.path.join(args.onnx, self.model_.target_name + ".onnx"),
+            providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
+        )
+        print("onnx load done")
+
     def demix(self, mix):
         samples = mix.shape[-1]
         margin = self.args.margin
-        chunk_size = self.args.chunks*44100
-        assert not margin == 0, 'margin cannot be zero!'
+        chunk_size = self.args.chunks * 44100
+        assert not margin == 0, "margin cannot be zero!"
         if margin > chunk_size:
             margin = chunk_size
 
         segmented_mix = {}
-        
+
         if self.args.chunks == 0 or samples < chunk_size:
             chunk_size = samples
-        
+
         counter = -1
         for skip in range(0, samples, chunk_size):
-            counter+=1
-    
+            counter += 1
+
             s_margin = 0 if counter == 0 else margin
-            end = min(skip+chunk_size+margin, samples)
+            end = min(skip + chunk_size + margin, samples)
 
-            start = skip-s_margin
+            start = skip - s_margin
 
-            segmented_mix[skip] = mix[:,start:end].copy()
+            segmented_mix[skip] = mix[:, start:end].copy()
             if end == samples:
                 break
 
         sources = self.demix_base(segmented_mix, margin_size=margin)
-        '''
+        """
         mix:(2,big_sample)
         segmented_mix:offset->(2,small_sample)
         sources:(1,2,big_sample)
-        '''
+        """
         return sources
+
     def demix_base(self, mixes, margin_size):
         chunked_sources = []
         progress_bar = tqdm(total=len(mixes))
@@ -106,84 +143,102 @@ def demix_base(self, mixes, margin_size):
             cmix = mixes[mix]
             sources = []
             n_sample = cmix.shape[1]
-            model=self.model_
-            trim = model.n_fft//2
-            gen_size = model.chunk_size-2*trim
-            pad = gen_size - n_sample%gen_size
-            mix_p = np.concatenate((np.zeros((2,trim)), cmix, np.zeros((2,pad)), np.zeros((2,trim))), 1)
+            model = self.model_
+            trim = model.n_fft // 2
+            gen_size = model.chunk_size - 2 * trim
+            pad = gen_size - n_sample % gen_size
+            mix_p = np.concatenate(
+                (np.zeros((2, trim)), cmix, np.zeros((2, pad)), np.zeros((2, trim))), 1
+            )
             mix_waves = []
             i = 0
             while i < n_sample + pad:
-                waves = np.array(mix_p[:, i:i+model.chunk_size])
+                waves = np.array(mix_p[:, i : i + model.chunk_size])
                 mix_waves.append(waves)
                 i += gen_size
             mix_waves = torch.tensor(mix_waves, dtype=torch.float32).to(cpu)
             with torch.no_grad():
                 _ort = self.model
                 spek = model.stft(mix_waves)
                 if self.args.denoise:
-                    spec_pred = -_ort.run(None, {'input': -spek.cpu().numpy()})[0]*0.5+_ort.run(None, {'input': spek.cpu().numpy()})[0]*0.5
+                    spec_pred = (
+                        -_ort.run(None, {"input": -spek.cpu().numpy()})[0] * 0.5
+                        + _ort.run(None, {"input": spek.cpu().numpy()})[0] * 0.5
+                    )
                     tar_waves = model.istft(torch.tensor(spec_pred))
                 else:
-                    tar_waves = model.istft(torch.tensor(_ort.run(None, {'input': spek.cpu().numpy()})[0]))
-                tar_signal = tar_waves[:,:,trim:-trim].transpose(0,1).reshape(2, -1).numpy()[:, :-pad]
+                    tar_waves = model.istft(
+                        torch.tensor(_ort.run(None, {"input": spek.cpu().numpy()})[0])
+                    )
+                tar_signal = (
+                    tar_waves[:, :, trim:-trim]
+                    .transpose(0, 1)
+                    .reshape(2, -1)
+                    .numpy()[:, :-pad]
+                )
 
                 start = 0 if mix == 0 else margin_size
                 end = None if mix == list(mixes.keys())[::-1][0] else -margin_size
                 if margin_size == 0:
                     end = None
-                sources.append(tar_signal[:,start:end])
+                sources.append(tar_signal[:, start:end])
 
                 progress_bar.update(1)
-            
+
             chunked_sources.append(sources)
         _sources = np.concatenate(chunked_sources, axis=-1)
         # del self.model
         progress_bar.close()
         return _sources
-    def prediction(self, m,vocal_root,others_root,format):
-        os.makedirs(vocal_root,exist_ok=True)
-        os.makedirs(others_root,exist_ok=True)
+
+    def prediction(self, m, vocal_root, others_root, format):
+        os.makedirs(vocal_root, exist_ok=True)
+        os.makedirs(others_root, exist_ok=True)
         basename = os.path.basename(m)
         mix, rate = librosa.load(m, mono=False, sr=44100)
         if mix.ndim == 1:
-            mix = np.asfortranarray([mix,mix])
+            mix = np.asfortranarray([mix, mix])
         mix = mix.T
         sources = self.demix(mix.T)
-        opt=sources[0].T
-        sf.write("%s/%s_main_vocal.%s"%(vocal_root,basename,format), mix-opt, rate)
-        sf.write("%s/%s_others.%s"%(others_root,basename,format), opt , rate)
-
-class MDXNetDereverb():
-    def __init__(self,chunks):
-        self.onnx="uvr5_weights/onnx_dereverb_By_FoxJoy"
-        self.shifts=10#'Predict with randomised equivariant stabilisation'
-        self.mixing="min_mag"#['default','min_mag','max_mag']
-        self.chunks=chunks
-        self.margin=44100
-        self.dim_t=9
-        self.dim_f=3072
-        self.n_fft=6144
-        self.denoise=True
-        self.pred=Predictor(self)
-
-    def _path_audio_(self,input,vocal_root,others_root,format):
-        self.pred.prediction(input,vocal_root,others_root,format)
-
-if __name__ == '__main__':
-    dereverb=MDXNetDereverb(15)
+        opt = sources[0].T
+        sf.write(
+            "%s/%s_main_vocal.%s" % (vocal_root, basename, format), mix - opt, rate
+        )
+        sf.write("%s/%s_others.%s" % (others_root, basename, format), opt, rate)
+
+
+class MDXNetDereverb:
+    def __init__(self, chunks):
+        self.onnx = "uvr5_weights/onnx_dereverb_By_FoxJoy"
+        self.shifts = 10  #'Predict with randomised equivariant stabilisation'
+        self.mixing = "min_mag"  # ['default','min_mag','max_mag']
+        self.chunks = chunks
+        self.margin = 44100
+        self.dim_t = 9
+        self.dim_f = 3072
+        self.n_fft = 6144
+        self.denoise = True
+        self.pred = Predictor(self)
+
+    def _path_audio_(self, input, vocal_root, others_root, format):
+        self.pred.prediction(input, vocal_root, others_root, format)
+
+
+if __name__ == "__main__":
+    dereverb = MDXNetDereverb(15)
     from time import time as ttime
-    t0=ttime()
+
+    t0 = ttime()
     dereverb._path_audio_(
         "雪雪伴奏对消HP5.wav",
         "vocal",
         "others",
     )
-    t1=ttime()
-    print(t1-t0)
+    t1 = ttime()
+    print(t1 - t0)
 
 
-'''
+"""
 
 runtime\python.exe MDXNet.py 
 
@@ -195,4 +250,4 @@ def _path_audio_(self,input,vocal_root,others_root,format):
 half15:0.7G->6.6G,22.69s
 fp32-15:0.7G->6.6G,20.85s
 
-'''
+"""