diff --git a/tada/modules/aligner_test.py b/tada/modules/aligner_test.py index 5464d22..ed619d2 100644 --- a/tada/modules/aligner_test.py +++ b/tada/modules/aligner_test.py @@ -4,7 +4,7 @@ import torch import torchaudio -from ..utils.test_utils import get_sample_dir +from ..utils.test_utils import get_sample_dir, skip_if_hub_model_unavailable from .aligner import Aligner @@ -17,6 +17,7 @@ ) def test_aligner(model_name_or_path: str): device = "cpu" + skip_if_hub_model_unavailable(model_name_or_path, subfolder="aligner") aligner = Aligner.from_pretrained(model_name_or_path, subfolder="aligner").to(device) audio, sample_rate = torchaudio.load(os.path.join(get_sample_dir(), "ljspeech.wav")) audio = audio.to(device) diff --git a/tada/modules/decoder_test.py b/tada/modules/decoder_test.py index 96d2ce3..ba08871 100644 --- a/tada/modules/decoder_test.py +++ b/tada/modules/decoder_test.py @@ -4,7 +4,7 @@ import torch import torchaudio -from ..utils.test_utils import get_sample_dir +from ..utils.test_utils import get_sample_dir, skip_if_hub_model_unavailable from .decoder import Decoder, DecoderConfig from .encoder import Encoder, EncoderConfig @@ -23,6 +23,8 @@ def test_decoder(model_name_or_path: str | None): encoder = Encoder(EncoderConfig()) decoder = Decoder(DecoderConfig()) else: + skip_if_hub_model_unavailable(model_name_or_path, subfolder="encoder") + skip_if_hub_model_unavailable(model_name_or_path, subfolder="decoder") encoder = Encoder.from_pretrained(model_name_or_path, subfolder="encoder").to(device) decoder = Decoder.from_pretrained(model_name_or_path, subfolder="decoder").to(device) audio, sample_rate = torchaudio.load(os.path.join(get_sample_dir(), "ljspeech.wav")) diff --git a/tada/modules/encoder_test.py b/tada/modules/encoder_test.py index 0df06da..c2e6886 100644 --- a/tada/modules/encoder_test.py +++ b/tada/modules/encoder_test.py @@ -3,11 +3,12 @@ import torch import torchaudio -from ..utils.test_utils import get_sample_dir +from ..utils.test_utils import get_sample_dir, skip_if_hub_model_unavailable from .encoder import Encoder def test_encoder(model_name_or_path: str = "HumeAI/tada-codec", device: str = "cuda"): + skip_if_hub_model_unavailable(model_name_or_path, subfolder="encoder") encoder = Encoder.from_pretrained(model_name_or_path, subfolder="encoder").to(device) audio, sample_rate = torchaudio.load(os.path.join(get_sample_dir(), "ljspeech.wav")) audio = audio.to(device) @@ -22,6 +23,7 @@ def test_token_values_are_local(model_name_or_path: str = "HumeAI/tada-codec", d """Test that token values are local. A change in one audio sample should only affect the immediately preceding and following token values. """ + skip_if_hub_model_unavailable(model_name_or_path, subfolder="encoder") encoder = Encoder.from_pretrained(model_name_or_path, subfolder="encoder").to(device) encoder.eval() torch.manual_seed(42) diff --git a/tada/modules/tada_test.py b/tada/modules/tada_test.py index edc8483..66dd50e 100644 --- a/tada/modules/tada_test.py +++ b/tada/modules/tada_test.py @@ -4,7 +4,7 @@ import torch import torchaudio -from ..utils.test_utils import get_sample_dir +from ..utils.test_utils import get_sample_dir, skip_if_hub_model_unavailable from .encoder import Encoder, EncoderConfig from .tada import TadaConfig, TadaForCausalLM @@ -26,6 +26,8 @@ def test_tada_for_causal_lm(model_name_or_path: str | None): ) model = TadaForCausalLM(config) else: + skip_if_hub_model_unavailable(model_name_or_path, subfolder="encoder") + skip_if_hub_model_unavailable(model_name_or_path, subfolder="llm") encoder = Encoder.from_pretrained("HumeAI/TADA", subfolder="encoder").to(device) model = TadaForCausalLM.from_pretrained(model_name_or_path, subfolder="llm").to(device) encoder.to(device).eval() diff --git a/tada/utils/test_utils.py b/tada/utils/test_utils.py index 7304857..cb42ae2 100644 --- a/tada/utils/test_utils.py +++ b/tada/utils/test_utils.py @@ -1,6 +1,8 @@ import os +import pytest import torch +from huggingface_hub import hf_hub_download def get_sample_dir() -> str: @@ -18,3 +20,11 @@ def get_default_device() -> torch.device: return torch.device("mps") else: return torch.device("cpu") + + +def skip_if_hub_model_unavailable(repo_id: str, subfolder: str | None = None, filename: str = "config.json") -> None: + """Skip integration tests when model artifacts cannot be fetched from Hugging Face Hub.""" + try: + hf_hub_download(repo_id=repo_id, filename=filename, subfolder=subfolder) + except Exception as exc: # pragma: no cover - best effort network check + pytest.skip(f"Hugging Face model '{repo_id}' is unavailable in this environment: {exc}")