Merge pull request #5 from NielsRogge/add_hf

Improve HF integration
JusperLee · Sep 21, 2024 · e7ec6cd · e7ec6cd
2 parents 851a8c2 + e9d2942
commit e7ec6cd
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/look2hear/models/base_model.py b/look2hear/models/base_model.py
@@ -7,6 +7,8 @@
 import torch
 import torch.nn as nn
 
+from huggingface_hub import PyTorchModelHubMixin
+
 
 def _unsqueeze_to_3d(x):
     """Normalize shape of `x` to [batch, n_chan, time]."""
@@ -32,7 +34,7 @@ def pad_to_appropriate_length(x, lcm):
     return x
 
 
-class BaseModel(nn.Module):
+class BaseModel(nn.Module, PyTorchModelHubMixin, repo_url="https://github.com/JusperLee/Apollo", pipeline_tag="audio-to-audio"):
     def __init__(self, sample_rate, in_chan=1):
         super().__init__()
         self._sample_rate = sample_rate