[TTA Pipeline] Test MusicGen and VITS (#26146)

This commit is contained in:
Sanchit Gandhi 2023-09-15 10:00:36 +01:00 committed by GitHub
parent 869733ab62
commit d70fab8b20
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 4 additions and 2 deletions

View File

@ -502,7 +502,7 @@ class MusicgenTester:
class MusicgenTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin, unittest.TestCase):
all_model_classes = (MusicgenForConditionalGeneration,) if is_torch_available() else ()
greedy_sample_model_classes = (MusicgenForConditionalGeneration,) if is_torch_available() else ()
pipeline_model_mapping = {}
pipeline_model_mapping = {"text-to-audio": MusicgenForConditionalGeneration} if is_torch_available() else {}
test_pruning = False # training is not supported yet for MusicGen
test_headmasking = False
test_resize_embeddings = False

View File

@ -40,6 +40,7 @@ from ...test_modeling_common import (
ids_tensor,
random_attention_mask,
)
from ...test_pipeline_mixin import PipelineTesterMixin
if is_torch_available():
@ -153,8 +154,9 @@ class VitsModelTester:
@require_torch
class VitsModelTest(ModelTesterMixin, unittest.TestCase):
class VitsModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase):
all_model_classes = (VitsModel,) if is_torch_available() else ()
pipeline_model_mapping = {"text-to-audio": VitsModel} if is_torch_available() else {}
is_encoder_decoder = False
test_pruning = False
test_headmasking = False