Merge branch 'develop' into add-manual-optimization-to-core-task

CHANGELOG.md

@@ -2,6 +2,20 @@
 
 ## develop
 
+### TL;DR
+
+#### Quality of life improvements
+
+Models can now be stored alongside their pipelines in the same repository, streamlining gating mechanism:
+- accept `pyannote/speaker-diarization-x.x` pipeline user agreement
+- ~~accept `pyannote/segmentation-3.0` model user agreement~~
+- ~~accept `pyannote/wespeaker-voxceleb-resnet34-LM` model user agreement~~
+- load pipeline with `Pipeline.from_pretrained("pyannote/speaker-diarization-3.1", use_auth_token=True)`
+
+#### Improve speech separation quality
+
+Clipping and speaker/source alignment issues in speech separation pipeline have been fixed.
+
 ### Breaking changes
 
 - BREAKING(task): drop support for `multilabel` training in `SpeakerDiarization` task
@@ -11,6 +25,7 @@
 
 ### New features
 
+- improve(hub): add support for pipeline repos that also include underlying models
 - feat(clustering): add support for `k-means` clustering
 - feat(model): add `wav2vec_frozen` option to freeze/unfreeze `wav2vec` in `SSeRiouSS` architecture
 - feat(task): add support for manual optimization in `SpeakerDiarization` task

pyannote/audio/core/model.py

@@ -532,6 +532,7 @@ def from_pretrained(
         map_location=None,
         hparams_file: Union[Path, Text] = None,
         strict: bool = True,
+        subfolder: Optional[str] = None,
         use_auth_token: Union[Text, None] = None,
         cache_dir: Union[Path, Text] = CACHE_DIR,
         **kwargs,
@@ -542,7 +543,7 @@ def from_pretrained(
         ----------
         checkpoint : Path or str
             Path to checkpoint, or a remote URL, or a model identifier from
-            the huggingface.co model hub.
+            the hf.co model hub.
         map_location: optional
             Same role as in torch.load().
             Defaults to `lambda storage, loc: storage`.
@@ -559,8 +560,10 @@ def from_pretrained(
         strict : bool, optional
             Whether to strictly enforce that the keys in checkpoint match
             the keys returned by this module’s state dict. Defaults to True.
+        subfolder : str, optional
+            Folder inside the hf.co model repo.
         use_auth_token : str, optional
-            When loading a private huggingface.co model, set `use_auth_token`
+            When loading a private hf.co model, set `use_auth_token`
             to True or to a string containing your hugginface.co authentication
             token that can be obtained by running `huggingface-cli login`
         cache_dir: Path or str, optional
@@ -606,18 +609,13 @@ def from_pretrained(
                 path_for_pl = hf_hub_download(
                     model_id,
                     HF_PYTORCH_WEIGHTS_NAME,
+                    subfolder=subfolder,
                     repo_type="model",
                     revision=revision,
                     library_name="pyannote",
                     library_version=__version__,
                     cache_dir=cache_dir,
-                    # force_download=False,
-                    # proxies=None,
-                    # etag_timeout=10,
-                    # resume_download=False,
                     use_auth_token=use_auth_token,
-                    # local_files_only=False,
-                    # legacy_cache_layout=False,
                 )
             except RepositoryNotFoundError:
                 print(

pyannote/audio/core/pipeline.py

@@ -47,6 +47,49 @@
 PIPELINE_PARAMS_NAME = "config.yaml"
 
 
+def expand_subfolders(
+    config, hf_model_id, use_auth_token: Optional[Text] = None
+) -> None:
+    """Expand $model subfolders in config
+
+    Processes `config` dictionary recursively and replaces "$model/{subfolder}"
+    values with {"checkpoint": hf_model_id,
+                 "subfolder": {subfolder},
+                 "use_auth_token": use_auth_token}
+
+    Parameters
+    ----------
+    config : dict
+    hf_model_id : str
+        Parent Huggingface model identifier
+    use_auth_token : str, optional
+        Token used for loading from the root folder.
+    """
+    if isinstance(config, dict):
+        for key, value in config.items():
+            if isinstance(value, str) and value.startswith("$model/"):
+                subfolder = "/".join(value.split("/")[1:])
+                config[key] = {
+                    "checkpoint": hf_model_id,
+                    "subfolder": subfolder,
+                    "use_auth_token": use_auth_token,
+                }
+            else:
+                expand_subfolders(value, hf_model_id, use_auth_token=use_auth_token)
+
+    elif isinstance(config, list):
+        for idx, value in enumerate(config):
+            if isinstance(value, str) and value.startswith("$model/"):
+                subfolder = "/".join(value.split("/")[1:])
+                config[idx] = {
+                    "checkpoint": hf_model_id,
+                    "subfolder": subfolder,
+                    "use_auth_token": use_auth_token,
+                }
+            else:
+                expand_subfolders(value, hf_model_id, use_auth_token=use_auth_token)
+
+
 class Pipeline(_Pipeline):
     @classmethod
     def from_pretrained(
@@ -95,13 +138,7 @@ def from_pretrained(
                     library_name="pyannote",
                     library_version=__version__,
                     cache_dir=cache_dir,
-                    # force_download=False,
-                    # proxies=None,
-                    # etag_timeout=10,
-                    # resume_download=False,
                     use_auth_token=use_auth_token,
-                    # local_files_only=False,
-                    # legacy_cache_layout=False,
                 )
 
             except RepositoryNotFoundError:
@@ -122,6 +159,7 @@ def from_pretrained(
 
         with open(config_yml, "r") as fp:
             config = yaml.load(fp, Loader=yaml.SafeLoader)
+        expand_subfolders(config, model_id, use_auth_token=use_auth_token)
 
         if "version" in config:
             check_version(