diff --git a/tests/engine/test_load.py b/tests/engine/test_load.py index 271e52e6..aea9f724 100644 --- a/tests/engine/test_load.py +++ b/tests/engine/test_load.py @@ -11,7 +11,7 @@ @pytest.mark.parametrize( "model_name, expected_output, should_raise", [ - ("PrunaAI/test-load-tiny-stable-diffusion-pipe-smashed", "PrunaModel", False), + ("pruna-test/test-load-tiny-stable-diffusion-pipe-smashed", "PrunaModel", False), ("NonExistentRepo/model", None, True), ], ) diff --git a/tests/engine/test_save.py b/tests/engine/test_save.py index b4190629..2cd6f3e1 100644 --- a/tests/engine/test_save.py +++ b/tests/engine/test_save.py @@ -15,14 +15,12 @@ from pruna.engine.pruna_model import PrunaModel - -@pytest.mark.skipif("HF_TOKEN" not in os.environ, reason="HF_TOKEN environment variable is not set, skipping tests.") @pytest.mark.slow @pytest.mark.cpu def test_save_llm_to_hub() -> None: """Test saving an LLM model to the Hugging Face Hub.""" download_repo_id = "hf-internal-testing/tiny-random-llama4" - upload_repo_id = "PrunaAI/test-save-tiny-random-llama4-smashed" + upload_repo_id = "pruna-test/test-save-tiny-random-llama4-smashed" model = AutoModelForCausalLM.from_pretrained(download_repo_id) smash_config = SmashConfig(device="cpu") pruna_model = smash( @@ -31,13 +29,12 @@ def test_save_llm_to_hub() -> None: ) pruna_model.push_to_hub(upload_repo_id, private=False) -@pytest.mark.skipif("HF_TOKEN" not in os.environ, reason="HF_TOKEN environment variable is not set, skipping tests.") @pytest.mark.slow @pytest.mark.cpu def test_save_diffusers_to_hub() -> None: """Test saving a diffusers model to the Hugging Face Hub.""" download_repo_id = "hf-internal-testing/tiny-stable-diffusion-pipe" - upload_repo_id = "PrunaAI/test-save-tiny-stable-diffusion-pipe-smashed" + upload_repo_id = "pruna-test/test-save-tiny-stable-diffusion-pipe-smashed" model = DiffusionPipeline.from_pretrained(download_repo_id) smash_config = SmashConfig(device="cpu") pruna_model = smash( diff --git a/tests/fixtures.py b/tests/fixtures.py index ae800b04..4022b6c5 100644 --- a/tests/fixtures.py +++ b/tests/fixtures.py @@ -53,7 +53,7 @@ def dataloader_fixture(request: pytest.FixtureRequest) -> Any: def whisper_tiny_random_model() -> tuple[Any, SmashConfig]: """Whisper tiny random model for speech recognition.""" - model_id = "PrunaAI/whisper-v3-tiny-random" + model_id = "pruna-test/whisper-v3-tiny-random" model = pipeline( "automatic-speech-recognition", model=model_id, @@ -181,8 +181,8 @@ def get_automodel_image_text_to_text_transformers(model_id: str) -> tuple[Any, S ), "dummy_lambda": dummy_model, # image generation AR models - "tiny_janus_pro": partial(get_automodel_image_text_to_text_transformers, "loulou2/tiny_janus"), - "wan_tiny_random": partial(get_diffusers_model, "PrunaAI/wan-t2v-tiny-random", torch_dtype=torch.bfloat16), - "flux_tiny": partial(get_diffusers_model, "loulou2/tiny_flux", torch_dtype=torch.float16), - "tiny_llama": partial(get_automodel_transformers, "loulou2/tiny_llama", torch_dtype=torch.bfloat16), + "tiny_janus_pro": partial(get_automodel_image_text_to_text_transformers, "pruna-test/tiny_janus"), + "wan_tiny_random": partial(get_diffusers_model, "pruna-test/wan-t2v-tiny-random", torch_dtype=torch.bfloat16), + "flux_tiny": partial(get_diffusers_model, "pruna-test/tiny_flux", torch_dtype=torch.float16), + "tiny_llama": partial(get_automodel_transformers, "pruna-test/tiny_llama", torch_dtype=torch.bfloat16), }