chore: fix all linting issues

Author: acolombier

demucs/api.py

@@ -22,7 +22,7 @@
 
 import subprocess
 
-from . import audio_legacy
+from . import audio_legacy  # noqa: F401
 import torch as th
 import torchaudio as ta
 
@@ -118,9 +118,17 @@ def __init__(
         self._name = model
         self._repo = repo
         self._load_model()
-        self.update_parameter(device=device, shifts=shifts, overlap=overlap, split=split,
-                              segment=segment, jobs=jobs, progress=progress, callback=callback,
-                              callback_arg=callback_arg)
+        self.update_parameter(
+            device=device,
+            shifts=shifts,
+            overlap=overlap,
+            split=split,
+            segment=segment,
+            jobs=jobs,
+            progress=progress,
+            callback=callback,
+            callback_arg=callback_arg,
+        )
 
     def update_parameter(
         self,
@@ -131,9 +139,7 @@ def update_parameter(
         segment: Optional[Union[int, _NotProvided]] = NotProvided,
         jobs: Union[int, _NotProvided] = NotProvided,
         progress: Union[bool, _NotProvided] = NotProvided,
-        callback: Optional[
-            Union[Callable[[dict], None], _NotProvided]
-        ] = NotProvided,
+        callback: Optional[Union[Callable[[dict], None], _NotProvided]] = NotProvided,
         callback_arg: Optional[Union[dict, _NotProvided]] = NotProvided,
     ):
         """
@@ -213,8 +219,9 @@ def _load_audio(self, track: Path):
         wav = None
 
         try:
-            wav = AudioFile(track).read(streams=0, samplerate=self._samplerate,
-                                        channels=self._audio_channels)
+            wav = AudioFile(track).read(
+                streams=0, samplerate=self._samplerate, channels=self._audio_channels
+            )
         except FileNotFoundError:
             errors["ffmpeg"] = "FFmpeg is not installed."
         except subprocess.CalledProcessError:
@@ -269,20 +276,20 @@ def separate_tensor(
         wav -= ref.mean()
         wav /= ref.std() + 1e-8
         out = apply_model(
-                self._model,
-                wav[None],
-                segment=self._segment,
-                shifts=self._shifts,
-                split=self._split,
-                overlap=self._overlap,
-                device=self._device,
-                num_workers=self._jobs,
-                callback=self._callback,
-                callback_arg=_replace_dict(
-                    self._callback_arg, ("audio_length", wav.shape[1])
-                ),
-                progress=self._progress,
-            )
+            self._model,
+            wav[None],
+            segment=self._segment,
+            shifts=self._shifts,
+            split=self._split,
+            overlap=self._overlap,
+            device=self._device,
+            num_workers=self._jobs,
+            callback=self._callback,
+            callback_arg=_replace_dict(
+                self._callback_arg, ("audio_length", wav.shape[1])
+            ),
+            progress=self._progress,
+        )
         if out is None:
             raise KeyboardInterrupt
         out *= ref.std() + 1e-8
@@ -336,7 +343,7 @@ def list_models(repo: Optional[Path] = None) -> Dict[str, Dict[str, Union[str, P
     """
     model_repo: ModelOnlyRepo
     if repo is None:
-        models = _parse_remote_files(REMOTE_ROOT / 'files.txt')
+        models = _parse_remote_files(REMOTE_ROOT / "files.txt")
         model_repo = RemoteRepo(models)
         bag_repo = BagOnlyRepo(REMOTE_ROOT, model_repo)
     else:
@@ -363,7 +370,7 @@ def list_models(repo: Optional[Path] = None) -> Dict[str, Dict[str, Union[str, P
         split=args.split,
         segment=args.segment,
         jobs=args.jobs,
-        callback=print
+        callback=print,
     )
     out = args.out / args.name
     out.mkdir(parents=True, exist_ok=True)

demucs/apply.py

@@ -57,7 +57,7 @@ def __init__(self, models: tp.List[Model],
         self.audio_channels = first.audio_channels
         self.samplerate = first.samplerate
         self.sources = first.sources
-        self.models = nn.ModuleList(models)
+        self.models = tp.cast(tp.List[Model], nn.ModuleList(models))
 
         if weights is None:
             weights = [[1. for _ in first.sources] for _ in models]
@@ -142,7 +142,7 @@ def _replace_dict(_dict: tp.Optional[dict], *subs: tp.Tuple[tp.Hashable, tp.Any]
     return _dict
 
 
-def apply_model(model: tp.Union[BagOfModels, Model],
+def apply_model(model: tp.Union[BagOfModels, Demucs, HDemucs, HTDemucs],
                 mix: tp.Union[th.Tensor, TensorChunk],
                 shifts: int = 1, split: bool = True,
                 overlap: float = 0.25, transition_power: float = 1.,

demucs/audio.py

@@ -10,7 +10,7 @@
 import lameenc
 import julius
 import numpy as np
-from . import audio_legacy
+from . import audio_legacy  # noqa: F401
 import torch
 import torchaudio as ta
 import typing as tp
@@ -19,18 +19,27 @@
 
 
 def _read_info(path):
-    stdout_data = sp.check_output([
-        'ffprobe', "-loglevel", "panic",
-        str(path), '-print_format', 'json', '-show_format', '-show_streams'
-    ])
-    return json.loads(stdout_data.decode('utf-8'))
+    stdout_data = sp.check_output(
+        [
+            "ffprobe",
+            "-loglevel",
+            "panic",
+            str(path),
+            "-print_format",
+            "json",
+            "-show_format",
+            "-show_streams",
+        ]
+    )
+    return json.loads(stdout_data.decode("utf-8"))
 
 
 class AudioFile:
     """
     Allows to read audio from any format supported by ffmpeg, as well as resampling or
     converting to mono on the fly. See :method:`read` for more details.
     """
+
     def __init__(self, path: Path):
         self.path = Path(path)
         self._info = None
@@ -51,30 +60,33 @@ def info(self):
 
     @property
     def duration(self):
-        return float(self.info['format']['duration'])
+        return float(self.info["format"]["duration"])
 
     @property
     def _audio_streams(self):
         return [
-            index for index, stream in enumerate(self.info["streams"])
+            index
+            for index, stream in enumerate(self.info["streams"])
             if stream["codec_type"] == "audio"
         ]
 
     def __len__(self):
         return len(self._audio_streams)
 
     def channels(self, stream=0):
-        return int(self.info['streams'][self._audio_streams[stream]]['channels'])
+        return int(self.info["streams"][self._audio_streams[stream]]["channels"])
 
     def samplerate(self, stream=0):
-        return int(self.info['streams'][self._audio_streams[stream]]['sample_rate'])
-
-    def read(self,
-             seek_time=None,
-             duration=None,
-             streams=slice(None),
-             samplerate=None,
-             channels=None):
+        return int(self.info["streams"][self._audio_streams[stream]]["sample_rate"])
+
+    def read(
+        self,
+        seek_time=None,
+        duration=None,
+        streams=slice(None),
+        samplerate=None,
+        channels=None,
+    ):
         """
         Slightly more efficient implementation than stempeg,
         in particular, this will extract all stems at once
@@ -106,22 +118,24 @@ def read(self,
             query_duration = None
         else:
             target_size = int((samplerate or self.samplerate()) * duration)
-            query_duration = float((target_size + 1) / (samplerate or self.samplerate()))
+            query_duration = float(
+                (target_size + 1) / (samplerate or self.samplerate())
+            )
 
         with temp_filenames(len(streams)) as filenames:
-            command = ['ffmpeg', '-y']
-            command += ['-loglevel', 'panic']
+            command = ["ffmpeg", "-y"]
+            command += ["-loglevel", "panic"]
             if seek_time:
-                command += ['-ss', str(seek_time)]
-            command += ['-i', str(self.path)]
+                command += ["-ss", str(seek_time)]
+            command += ["-i", str(self.path)]
             for stream, filename in zip(streams, filenames):
-                command += ['-map', f'0:{self._audio_streams[stream]}']
+                command += ["-map", f"0:{self._audio_streams[stream]}"]
                 if query_duration is not None:
-                    command += ['-t', str(query_duration)]
-                command += ['-threads', '1']
-                command += ['-f', 'f32le']
+                    command += ["-t", str(query_duration)]
+                command += ["-threads", "1"]
+                command += ["-f", "f32le"]
                 if samplerate is not None:
-                    command += ['-ar', str(samplerate)]
+                    command += ["-ar", str(samplerate)]
                 command += [filename]
 
             sp.run(command, check=True)
@@ -163,7 +177,9 @@ def convert_audio_channels(wav, channels=2):
         wav = wav[..., :channels, :]
     else:
         # Case 4: What is a reasonable choice here?
-        raise ValueError('The audio file has less channels than requested but is not mono.')
+        raise ValueError(
+            "The audio file has less channels than requested but is not mono."
+        )
     return wav
 
 
@@ -216,32 +232,34 @@ def encode_mp3(wav, path, samplerate=44100, bitrate=320, quality=2, verbose=Fals
         f.write(mp3_data)
 
 
-def prevent_clip(wav, mode='rescale'):
+def prevent_clip(wav, mode="rescale"):
     """
     different strategies for avoiding raw clipping.
     """
-    if mode is None or mode == 'none':
+    if mode is None or mode == "none":
         return wav
     assert wav.dtype.is_floating_point, "too late for clipping"
-    if mode == 'rescale':
+    if mode == "rescale":
         wav = wav / max(1.01 * wav.abs().max(), 1)
-    elif mode == 'clamp':
+    elif mode == "clamp":
         wav = wav.clamp(-0.99, 0.99)
-    elif mode == 'tanh':
+    elif mode == "tanh":
         wav = torch.tanh(wav)
     else:
         raise ValueError(f"Invalid mode {mode}")
     return wav
 
 
-def save_audio(wav: torch.Tensor,
-               path: tp.Union[str, Path],
-               samplerate: int,
-               bitrate: int = 320,
-               clip: tp.Literal["rescale", "clamp", "tanh", "none"] = 'rescale',
-               bits_per_sample: tp.Literal[16, 24, 32] = 16,
-               as_float: bool = False,
-               preset: tp.Literal[2, 3, 4, 5, 6, 7] = 2):
+def save_audio(
+    wav: torch.Tensor,
+    path: tp.Union[str, Path],
+    samplerate: int,
+    bitrate: int = 320,
+    clip: tp.Literal["rescale", "clamp", "tanh", "none"] = "rescale",
+    bits_per_sample: tp.Literal[16, 24, 32] = 16,
+    as_float: bool = False,
+    preset: tp.Literal[2, 3, 4, 5, 6, 7] = 2,
+):
     """Save audio file, automatically preventing clipping if necessary
     based on the given `clip` strategy. If the path ends in `.mp3`, this
     will save as mp3 with the given `bitrate`. Use `preset` to set mp3 quality:
@@ -255,11 +273,16 @@ def save_audio(wav: torch.Tensor,
     elif suffix == ".wav":
         if as_float:
             bits_per_sample = 32
-            encoding = 'PCM_F'
+            encoding = "PCM_F"
         else:
-            encoding = 'PCM_S'
-        ta.save(str(path), wav, sample_rate=samplerate,
-                encoding=encoding, bits_per_sample=bits_per_sample)
+            encoding = "PCM_S"
+        ta.save(
+            str(path),
+            wav,
+            sample_rate=samplerate,
+            encoding=encoding,
+            bits_per_sample=bits_per_sample,
+        )
     elif suffix == ".flac":
         ta.save(str(path), wav, sample_rate=samplerate, bits_per_sample=bits_per_sample)
     else:

demucs/audio_legacy.py

@@ -5,7 +5,7 @@
 import sys
 import warnings
 
-if not "torchaudio" in sys.modules:
+if "torchaudio" not in sys.modules:
     os.environ["TORCHAUDIO_USE_BACKEND_DISPATCHER"] = "0"
 elif os.getenv("TORCHAUDIO_USE_BACKEND_DISPATCHER", default="1") == "1":
     if sys.modules["torchaudio"].__version__ >= "2.1":

demucs/demucs.py

@@ -134,7 +134,17 @@ def __init__(self, channels: int, compress: float = 4, depth: int = 2, init: flo
         for d in range(self.depth):
             dilation = 2 ** d if dilate else 1
             padding = dilation * (kernel // 2)
-            mods = [
+            mods: tp.List[
+                nn.Conv1d |
+                nn.GroupNorm |
+                nn.Identity |
+                nn.GELU |
+                nn.ReLU |
+                nn.GLU |
+                LayerScale |
+                BLSTM |
+                LocalState
+            ] = [
                 nn.Conv1d(channels, hidden, kernel, dilation=dilation, padding=padding),
                 norm_fn(hidden), act(),
                 nn.Conv1d(hidden, 2 * channels, 1),