[CI] Update test fetcher script

.github/workflows/pr_test_fetcher.yml

@@ -28,7 +28,6 @@ jobs:
         shell: bash
     outputs:
       matrix: ${{ steps.set_matrix.outputs.matrix }}
-      test_map: ${{ steps.set_matrix.outputs.test_map }}
     steps:
     - name: Checkout diffusers
       uses: actions/checkout@v6
@@ -51,31 +50,26 @@ jobs:
         path: test_preparation.txt
     - id: set_matrix
       name: Create Test Matrix
-      # The `keys` is used as GitHub actions matrix for jobs, i.e. `models`, `pipelines`, etc.
-      # The `test_map` is used to get the actual identified test files under each key.
-      # If no test to run (so no `test_map.json` file), create a dummy map (empty matrix will fail)
+      # `test_map.json` is a list of matrix entries `{"name", "paths", "markers"}`, one job each.
+      # If no test to run (so no `test_map.json` file), emit an empty list and skip the test job.
       run: |
         if [ -f test_map.json ]; then
-            keys=$(python3 -c 'import json; fp = open("test_map.json"); test_map = json.load(fp); fp.close(); d = list(test_map.keys()); print(json.dumps(d))')
-            test_map=$(python3 -c 'import json; fp = open("test_map.json"); test_map = json.load(fp); fp.close(); print(json.dumps(test_map))')
+            matrix=$(python3 -c 'import json; print(json.dumps(json.load(open("test_map.json"))))')
         else
-            keys=$(python3 -c 'keys = ["dummy"]; print(keys)')
-            test_map=$(python3 -c 'test_map = {"dummy": []}; print(test_map)')
+            matrix='[]'
         fi
-        echo $keys
-        echo $test_map
-        echo "matrix=$keys" >> $GITHUB_OUTPUT
-        echo "test_map=$test_map" >> $GITHUB_OUTPUT
+        echo "matrix=$matrix" >> $GITHUB_OUTPUT
+        echo "Matrix: $matrix"
 
   run_pr_tests:
-    name: Run PR Tests
+    name: ${{ matrix.name }}
     needs: setup_pr_tests
-    if: contains(fromJson(needs.setup_pr_tests.outputs.matrix), 'dummy') != true
+    if: needs.setup_pr_tests.outputs.matrix != '[]'
     strategy:
       fail-fast: false
       max-parallel: 2
       matrix:
-        modules: ${{ fromJson(needs.setup_pr_tests.outputs.matrix) }}
+        include: ${{ fromJson(needs.setup_pr_tests.outputs.matrix) }}
     runs-on:
       group: aws-general-8-plus
     container:
@@ -101,20 +95,20 @@ jobs:
 
     - name: Run all selected tests on CPU
       run: |
-        pytest -n 2 --dist=loadfile -v --make-reports=${{ matrix.modules }}_tests_cpu ${{ fromJson(needs.setup_pr_tests.outputs.test_map)[matrix.modules] }}
+        pytest -n 2 --dist=loadfile -v -m "${{ matrix.markers }}" --make-reports=${{ matrix.name }}_tests_cpu ${{ matrix.paths }}
 
     - name: Failure short reports
       if: ${{ failure() }}
       continue-on-error: true
       run: |
-        cat reports/${{ matrix.modules }}_tests_cpu_stats.txt
-        cat reports/${{ matrix.modules }}_tests_cpu_failures_short.txt
+        cat reports/${{ matrix.name }}_tests_cpu_stats.txt
+        cat reports/${{ matrix.name }}_tests_cpu_failures_short.txt
 
     - name: Test suite reports artifacts
       if: ${{ always() }}
       uses: actions/upload-artifact@v6
       with:
-          name: ${{ matrix.modules }}_test_reports
+          name: ${{ matrix.name }}_test_reports
           path: reports
 
   run_staging_tests:

tests/lora/test_lora_layers_sd.py

@@ -24,8 +24,6 @@
 from transformers import CLIPTextModel, CLIPTokenizer
 
 from diffusers import (
-    AutoPipelineForImage2Image,
-    AutoPipelineForText2Image,
     DDIMScheduler,
     DiffusionPipeline,
     LCMScheduler,
@@ -641,6 +639,8 @@ def test_load_unload_load_kohya_lora(self):
         release_memory(pipe)
 
     def test_not_empty_state_dict(self):
+        from diffusers import AutoPipelineForText2Image
+
         # Makes sure https://github.com/huggingface/diffusers/issues/7054 does not happen again
         pipe = AutoPipelineForText2Image.from_pretrained(
             "stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16
@@ -655,6 +655,8 @@ def test_not_empty_state_dict(self):
         release_memory(pipe)
 
     def test_load_unload_load_state_dict(self):
+        from diffusers import AutoPipelineForText2Image
+
         # Makes sure https://github.com/huggingface/diffusers/issues/7054 does not happen again
         pipe = AutoPipelineForText2Image.from_pretrained(
             "stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16
@@ -705,6 +707,8 @@ def test_sdv1_5_lcm_lora(self):
         release_memory(pipe)
 
     def test_sdv1_5_lcm_lora_img2img(self):
+        from diffusers import AutoPipelineForImage2Image
+
         pipe = AutoPipelineForImage2Image.from_pretrained(
             "stable-diffusion-v1-5/stable-diffusion-v1-5", torch_dtype=torch.float16
         )

tests/models/test_modeling_common.py

@@ -26,7 +26,7 @@
 from huggingface_hub import ModelCard, delete_repo, snapshot_download, try_to_load_from_cache
 from huggingface_hub.utils import HfHubHTTPError, is_jinja_available
 
-from diffusers.models import FluxTransformer2DModel, SD3Transformer2DModel, UNet2DConditionModel
+from diffusers.models import UNet2DConditionModel
 
 from ..others.test_utils import TOKEN, USER, is_staging_test
 from ..testing_utils import (
@@ -118,6 +118,8 @@ def test_cached_files_are_used_when_no_internet(self):
                 assert False, "Parameters not the same!"
 
     def test_local_files_only_with_sharded_checkpoint(self):
+        from diffusers.models import FluxTransformer2DModel
+
         repo_id = "hf-internal-testing/tiny-flux-sharded"
         error_response = mock.Mock(
             status_code=500,
@@ -232,6 +234,8 @@ def test_keep_modules_in_fp32(self):
         A simple tests to check if the modules under `_keep_in_fp32_modules` are kept in fp32 when we load the model in fp16/bf16
         Also ensures if inference works.
         """
+        from diffusers.models import SD3Transformer2DModel
+
         fp32_modules = SD3Transformer2DModel._keep_in_fp32_modules
 
         for torch_dtype in [torch.bfloat16, torch.float16]:

tests/others/test_flashpack.py

@@ -15,7 +15,6 @@
 
 import pytest
 
-from diffusers import AutoPipelineForText2Image
 from diffusers.models.auto_model import AutoModel
 
 from ..testing_utils import is_torch_available, require_flashpack, require_torch_gpu
@@ -40,6 +39,8 @@ def test_save_load_model(self, tmp_path):
 
     @require_flashpack
     def test_save_load_pipeline(self, tmp_path):
+        from diffusers import AutoPipelineForText2Image
+
         pipeline = AutoPipelineForText2Image.from_pretrained(self.model_id)
         pipeline.save_pretrained(tmp_path, use_flashpack=True)
         assert (tmp_path / "transformer" / "model.flashpack").exists()

tests/pipelines/kandinsky3/test_kandinsky3.py

@@ -22,8 +22,6 @@
 from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
 
 from diffusers import (
-    AutoPipelineForImage2Image,
-    AutoPipelineForText2Image,
     Kandinsky3Pipeline,
     Kandinsky3UNet,
     VQModel,
@@ -185,6 +183,8 @@ def tearDown(self):
         backend_empty_cache(torch_device)
 
     def test_kandinskyV3(self):
+        from diffusers import AutoPipelineForText2Image
+
         pipe = AutoPipelineForText2Image.from_pretrained(
             "kandinsky-community/kandinsky-3", variant="fp16", torch_dtype=torch.float16
         )
@@ -211,6 +211,8 @@ def test_kandinskyV3(self):
         self.assertTrue(np.allclose(image_np, expected_image_np, atol=5e-2))
 
     def test_kandinskyV3_img2img(self):
+        from diffusers import AutoPipelineForImage2Image
+
         pipe = AutoPipelineForImage2Image.from_pretrained(
             "kandinsky-community/kandinsky-3", variant="fp16", torch_dtype=torch.float16
         )

tests/pipelines/kandinsky3/test_kandinsky3_img2img.py

@@ -23,7 +23,6 @@
 from transformers import AutoConfig, AutoTokenizer, T5EncoderModel
 
 from diffusers import (
-    AutoPipelineForImage2Image,
     Kandinsky3Img2ImgPipeline,
     Kandinsky3UNet,
     VQModel,
@@ -207,6 +206,8 @@ def tearDown(self):
         backend_empty_cache(torch_device)
 
     def test_kandinskyV3_img2img(self):
+        from diffusers import AutoPipelineForImage2Image
+
         pipe = AutoPipelineForImage2Image.from_pretrained(
             "kandinsky-community/kandinsky-3", variant="fp16", torch_dtype=torch.float16
         )

tests/pipelines/pag/test_pag_sd.py

@@ -23,7 +23,6 @@
 
 from diffusers import (
     AutoencoderKL,
-    AutoPipelineForText2Image,
     DDIMScheduler,
     StableDiffusionPAGPipeline,
     StableDiffusionPipeline,
@@ -315,6 +314,8 @@ def get_inputs(self, device, generator_device="cpu", seed=1, guidance_scale=7.0)
         return inputs
 
     def test_pag_cfg(self):
+        from diffusers import AutoPipelineForText2Image
+
         pipeline = AutoPipelineForText2Image.from_pretrained(self.repo_id, enable_pag=True, torch_dtype=torch.float16)
         pipeline.enable_model_cpu_offload(device=torch_device)
         pipeline.set_progress_bar_config(disable=None)
@@ -333,6 +334,8 @@ def test_pag_cfg(self):
         )
 
     def test_pag_uncond(self):
+        from diffusers import AutoPipelineForText2Image
+
         pipeline = AutoPipelineForText2Image.from_pretrained(self.repo_id, enable_pag=True, torch_dtype=torch.float16)
         pipeline.enable_model_cpu_offload(device=torch_device)
         pipeline.set_progress_bar_config(disable=None)

tests/pipelines/pag/test_pag_sd3_img2img.py

@@ -16,7 +16,6 @@
 
 from diffusers import (
     AutoencoderKL,
-    AutoPipelineForImage2Image,
     FlowMatchEulerDiscreteScheduler,
     SD3Transformer2DModel,
     StableDiffusion3Img2ImgPipeline,
@@ -240,6 +239,8 @@ def get_inputs(
         return inputs
 
     def test_pag_cfg(self):
+        from diffusers import AutoPipelineForImage2Image
+
         pipeline = AutoPipelineForImage2Image.from_pretrained(
             self.repo_id, enable_pag=True, torch_dtype=torch.float16, pag_applied_layers=["blocks.17"]
         )
@@ -268,6 +269,8 @@ def test_pag_cfg(self):
         )
 
     def test_pag_uncond(self):
+        from diffusers import AutoPipelineForImage2Image
+
         pipeline = AutoPipelineForImage2Image.from_pretrained(
             self.repo_id, enable_pag=True, torch_dtype=torch.float16, pag_applied_layers=["blocks.(4|17)"]
         )

tests/pipelines/pag/test_pag_sd_img2img.py

@@ -25,7 +25,6 @@
 from diffusers import (
     AutoencoderKL,
     AutoencoderTiny,
-    AutoPipelineForImage2Image,
     EulerDiscreteScheduler,
     StableDiffusionImg2ImgPipeline,
     StableDiffusionPAGImg2ImgPipeline,
@@ -254,6 +253,8 @@ def get_inputs(self, device, generator_device="cpu", dtype=torch.float32, seed=0
         return inputs
 
     def test_pag_cfg(self):
+        from diffusers import AutoPipelineForImage2Image
+
         pipeline = AutoPipelineForImage2Image.from_pretrained(self.repo_id, enable_pag=True, torch_dtype=torch.float16)
         pipeline.enable_model_cpu_offload(device=torch_device)
         pipeline.set_progress_bar_config(disable=None)
@@ -272,6 +273,8 @@ def test_pag_cfg(self):
         )
 
     def test_pag_uncond(self):
+        from diffusers import AutoPipelineForImage2Image
+
         pipeline = AutoPipelineForImage2Image.from_pretrained(self.repo_id, enable_pag=True, torch_dtype=torch.float16)
         pipeline.enable_model_cpu_offload(device=torch_device)
         pipeline.set_progress_bar_config(disable=None)

tests/pipelines/pag/test_pag_sd_inpaint.py

@@ -24,7 +24,6 @@
 
 from diffusers import (
     AutoencoderKL,
-    AutoPipelineForInpainting,
     PNDMScheduler,
     StableDiffusionPAGInpaintPipeline,
     UNet2DConditionModel,
@@ -289,6 +288,8 @@ def get_inputs(self, device, generator_device="cpu", seed=0, guidance_scale=7.0)
         return inputs
 
     def test_pag_cfg(self):
+        from diffusers import AutoPipelineForInpainting
+
         pipeline = AutoPipelineForInpainting.from_pretrained(self.repo_id, enable_pag=True, torch_dtype=torch.float16)
         pipeline.enable_model_cpu_offload(device=torch_device)
         pipeline.set_progress_bar_config(disable=None)
@@ -307,6 +308,8 @@ def test_pag_cfg(self):
         )
 
     def test_pag_uncond(self):
+        from diffusers import AutoPipelineForInpainting
+
         pipeline = AutoPipelineForInpainting.from_pretrained(self.repo_id, enable_pag=True, torch_dtype=torch.float16)
         pipeline.enable_model_cpu_offload(device=torch_device)
         pipeline.set_progress_bar_config(disable=None)

tests/pipelines/pag/test_pag_sdxl.py

@@ -23,7 +23,6 @@
 
 from diffusers import (
     AutoencoderKL,
-    AutoPipelineForText2Image,
     EulerDiscreteScheduler,
     StableDiffusionXLPAGPipeline,
     StableDiffusionXLPipeline,
@@ -320,6 +319,8 @@ def get_inputs(self, device, generator_device="cpu", seed=0, guidance_scale=7.0)
         return inputs
 
     def test_pag_cfg(self):
+        from diffusers import AutoPipelineForText2Image
+
         pipeline = AutoPipelineForText2Image.from_pretrained(self.repo_id, enable_pag=True, torch_dtype=torch.float16)
         pipeline.enable_model_cpu_offload(device=torch_device)
         pipeline.set_progress_bar_config(disable=None)
@@ -342,6 +343,8 @@ def test_pag_cfg(self):
         )
 
     def test_pag_uncond(self):
+        from diffusers import AutoPipelineForText2Image
+
         pipeline = AutoPipelineForText2Image.from_pretrained(self.repo_id, enable_pag=True, torch_dtype=torch.float16)
         pipeline.enable_model_cpu_offload(device=torch_device)
         pipeline.set_progress_bar_config(disable=None)

tests/pipelines/pag/test_pag_sdxl_img2img.py

@@ -32,7 +32,6 @@
 
 from diffusers import (
     AutoencoderKL,
-    AutoPipelineForImage2Image,
     EulerDiscreteScheduler,
     StableDiffusionXLImg2ImgPipeline,
     StableDiffusionXLPAGImg2ImgPipeline,
@@ -302,6 +301,8 @@ def get_inputs(self, device, generator_device="cpu", seed=0, guidance_scale=7.0)
         return inputs
 
     def test_pag_cfg(self):
+        from diffusers import AutoPipelineForImage2Image
+
         pipeline = AutoPipelineForImage2Image.from_pretrained(self.repo_id, enable_pag=True, torch_dtype=torch.float16)
         pipeline.enable_model_cpu_offload(device=torch_device)
         pipeline.set_progress_bar_config(disable=None)
@@ -319,6 +320,8 @@ def test_pag_cfg(self):
         )
 
     def test_pag_uncond(self):
+        from diffusers import AutoPipelineForImage2Image
+
         pipeline = AutoPipelineForImage2Image.from_pretrained(self.repo_id, enable_pag=True, torch_dtype=torch.float16)
         pipeline.enable_model_cpu_offload(device=torch_device)
         pipeline.set_progress_bar_config(disable=None)

tests/pipelines/pag/test_pag_sdxl_inpaint.py

@@ -33,7 +33,6 @@
 
 from diffusers import (
     AutoencoderKL,
-    AutoPipelineForInpainting,
     EulerDiscreteScheduler,
     StableDiffusionXLInpaintPipeline,
     StableDiffusionXLPAGInpaintPipeline,
@@ -308,6 +307,8 @@ def get_inputs(self, device, generator_device="cpu", seed=0, guidance_scale=7.0)
         return inputs
 
     def test_pag_cfg(self):
+        from diffusers import AutoPipelineForInpainting
+
         pipeline = AutoPipelineForInpainting.from_pretrained(self.repo_id, enable_pag=True, torch_dtype=torch.float16)
         pipeline.enable_model_cpu_offload(device=torch_device)
         pipeline.set_progress_bar_config(disable=None)
@@ -325,6 +326,8 @@ def test_pag_cfg(self):
         )
 
     def test_pag_uncond(self):
+        from diffusers import AutoPipelineForInpainting
+
         pipeline = AutoPipelineForInpainting.from_pretrained(self.repo_id, enable_pag=True, torch_dtype=torch.float16)
         pipeline.enable_model_cpu_offload(device=torch_device)
         pipeline.set_progress_bar_config(disable=None)