From b937ac8578992c02f86144446d968a04359be3c9 Mon Sep 17 00:00:00 2001 From: Jan Cuhel Date: Fri, 12 Jan 2024 23:57:20 +0100 Subject: [PATCH] Update --- tests/integration/test_pipeline.py | 1 - tests/unittests/test_image_generation.py | 10 +++++----- 2 files changed, 5 insertions(+), 6 deletions(-) diff --git a/tests/integration/test_pipeline.py b/tests/integration/test_pipeline.py index f236472..9ece194 100644 --- a/tests/integration/test_pipeline.py +++ b/tests/integration/test_pipeline.py @@ -378,7 +378,6 @@ def test_cuda_simple_sdxl_turbo_classification_pipeline(): _check_detection_pipeline(cmd, target_folder) - # def test_cpu_simple_sdxl_classification_pipeline(): # # Define target folder # target_folder = "data/data-cls-cpu-simple-sdxl/" diff --git a/tests/unittests/test_image_generation.py b/tests/unittests/test_image_generation.py index 8e044ad..e3cc9cf 100644 --- a/tests/unittests/test_image_generation.py +++ b/tests/unittests/test_image_generation.py @@ -50,13 +50,13 @@ def _check_image_generator(image_generator_class: Type[Union[StableDiffusionImag image_generator.release(empty_cuda_cache=True if device != "cpu" else False) -# @pytest.mark.skipif(not torch.cuda.is_available(), reason="Test requires GPU") -# def test_cuda_sdxl_image_generator(): -# _check_image_generator(StableDiffusionImageGenerator, "cuda") +@pytest.mark.skipif(not torch.cuda.is_available(), reason="Test requires GPU") +def test_cuda_sdxl_image_generator(): + _check_image_generator(StableDiffusionImageGenerator, "cuda") -# def test_cpu_sdxl_image_generator(): -# _check_image_generator(StableDiffusionImageGenerator, "cpu") +def test_cpu_sdxl_image_generator(): + _check_image_generator(StableDiffusionImageGenerator, "cpu") @pytest.mark.skipif(not torch.cuda.is_available(), reason="Test requires GPU")