test num images

huggingface · Nov 4, 2024 · a77549c · a77549c
1 parent d22382c
commit a77549c
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 9 deletions.
diff --git a/.github/workflows/test_onnxruntime.yml b/.github/workflows/test_onnxruntime.yml
@@ -39,18 +39,18 @@ jobs:
         run: |
           pip install --upgrade pip
           pip install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
-          pip install .[tests,onnxruntime]
+          pip install .[tests,onnxruntime] transformers[testing]
 
       - if: ${{ matrix.transformers-version != 'latest' }}
         name: Downgrade Transformers
         run: pip install transformers==${{ matrix.transformers-version }}
 
       - name: Test with pytest (in series)
         run: |
-          pytest tests/onnxruntime -m "run_in_series" --durations=0 -vvvv -s
+          pytest tests/onnxruntime -m "run_in_series" --durations=0
 
       - name: Test with pytest (in parallel)
         run: |
-          pytest tests/onnxruntime -m "not run_in_series" --durations=0 -vvvv -s -n auto
+          pytest tests/onnxruntime -m "not run_in_series" --durations=0 -n auto -k "test_num_images_per_prompt"
         env:
           HF_HUB_READ_TOKEN: ${{ secrets.HF_HUB_READ_TOKEN }}
diff --git a/tests/onnxruntime/test_diffusion.py b/tests/onnxruntime/test_diffusion.py
@@ -138,8 +138,8 @@ def test_num_images_per_prompt(self, model_arch: str):
         pipeline = self.ORTMODEL_CLASS.from_pretrained(self.onnx_model_dirs[model_arch])
 
         for batch_size in [1, 3]:
-            for height in [64, 128]:
-                for width in [64, 128]:
+            for height in [16, 32]:
+                for width in [16, 32]:
                     for num_images_per_prompt in [1, 3]:
                         inputs = self.generate_inputs(height=height, width=width, batch_size=batch_size)
                         outputs = pipeline(**inputs, num_images_per_prompt=num_images_per_prompt).images
@@ -379,8 +379,8 @@ def test_num_images_per_prompt(self, model_arch: str):
         pipeline = self.ORTMODEL_CLASS.from_pretrained(self.onnx_model_dirs[model_arch])
 
         for batch_size in [1, 3]:
-            for height in [64, 128]:
-                for width in [64, 128]:
+            for height in [16, 32]:
+                for width in [16, 32]:
                     for num_images_per_prompt in [1, 3]:
                         inputs = self.generate_inputs(height=height, width=width, batch_size=batch_size)
                         outputs = pipeline(**inputs, num_images_per_prompt=num_images_per_prompt).images
@@ -611,8 +611,8 @@ def test_num_images_per_prompt(self, model_arch: str):
         pipeline = self.ORTMODEL_CLASS.from_pretrained(self.onnx_model_dirs[model_arch])
 
         for batch_size in [1, 3]:
-            for height in [64, 128]:
-                for width in [64, 128]:
+            for height in [16, 32]:
+                for width in [16, 32]:
                     for num_images_per_prompt in [1, 3]:
                         inputs = self.generate_inputs(height=height, width=width, batch_size=batch_size)
                         outputs = pipeline(**inputs, num_images_per_prompt=num_images_per_prompt).images