diff --git a/.github/workflows/pr_modular_tests.yml b/.github/workflows/pr_modular_tests.yml new file mode 100644 index 0000000000..e01345e325 --- /dev/null +++ b/.github/workflows/pr_modular_tests.yml @@ -0,0 +1,141 @@ +name: Fast PR tests for Modular + +on: + pull_request: + branches: [main] + paths: + - "src/diffusers/modular_pipelines/**.py" + - "src/diffusers/models/modeling_utils.py" + - "src/diffusers/models/model_loading_utils.py" + - "src/diffusers/pipelines/pipeline_utils.py" + - "src/diffusers/pipeline_loading_utils.py" + - "src/diffusers/loaders/lora_base.py" + - "src/diffusers/loaders/lora_pipeline.py" + - "src/diffusers/loaders/peft.py" + - "tests/modular_pipelines/**.py" + - ".github/**.yml" + - "utils/**.py" + - "setup.py" + push: + branches: + - ci-* + +concurrency: + group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }} + cancel-in-progress: true + +env: + DIFFUSERS_IS_CI: yes + HF_HUB_ENABLE_HF_TRANSFER: 1 + OMP_NUM_THREADS: 4 + MKL_NUM_THREADS: 4 + PYTEST_TIMEOUT: 60 + +jobs: + check_code_quality: + runs-on: ubuntu-22.04 + steps: + - uses: actions/checkout@v3 + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: "3.10" + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install .[quality] + - name: Check quality + run: make quality + - name: Check if failure + if: ${{ failure() }} + run: | + echo "Quality check failed. Please ensure the right dependency versions are installed with 'pip install -e .[quality]' and run 'make style && make quality'" >> $GITHUB_STEP_SUMMARY + + check_repository_consistency: + needs: check_code_quality + runs-on: ubuntu-22.04 + steps: + - uses: actions/checkout@v3 + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: "3.10" + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install .[quality] + - name: Check repo consistency + run: | + python utils/check_copies.py + python utils/check_dummies.py + python utils/check_support_list.py + make deps_table_check_updated + - name: Check if failure + if: ${{ failure() }} + run: | + echo "Repo consistency check failed. Please ensure the right dependency versions are installed with 'pip install -e .[quality]' and run 'make fix-copies'" >> $GITHUB_STEP_SUMMARY + + run_fast_tests: + needs: [check_code_quality, check_repository_consistency] + strategy: + fail-fast: false + matrix: + config: + - name: Fast PyTorch Modular Pipeline CPU tests + framework: pytorch_pipelines + runner: aws-highmemory-32-plus + image: diffusers/diffusers-pytorch-cpu + report: torch_cpu_modular_pipelines + + name: ${{ matrix.config.name }} + + runs-on: + group: ${{ matrix.config.runner }} + + container: + image: ${{ matrix.config.image }} + options: --shm-size "16gb" --ipc host -v /mnt/hf_cache:/mnt/cache/ + + defaults: + run: + shell: bash + + steps: + - name: Checkout diffusers + uses: actions/checkout@v3 + with: + fetch-depth: 2 + + - name: Install dependencies + run: | + python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" + python -m uv pip install -e [quality,test] + pip uninstall transformers -y && python -m uv pip install -U transformers@git+https://github.com/huggingface/transformers.git --no-deps + pip uninstall accelerate -y && python -m uv pip install -U accelerate@git+https://github.com/huggingface/accelerate.git --no-deps + + - name: Environment + run: | + python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" + python utils/print_env.py + + - name: Run fast PyTorch Pipeline CPU tests + if: ${{ matrix.config.framework == 'pytorch_pipelines' }} + run: | + python -m venv /opt/venv && export PATH="/opt/venv/bin:$PATH" + python -m pytest -n 8 --max-worker-restart=0 --dist=loadfile \ + -s -v -k "not Flax and not Onnx" \ + --make-reports=tests_${{ matrix.config.report }} \ + tests/modular_pipelines + + - name: Failure short reports + if: ${{ failure() }} + run: cat reports/tests_${{ matrix.config.report }}_failures_short.txt + + - name: Test suite reports artifacts + if: ${{ always() }} + uses: actions/upload-artifact@v4 + with: + name: pr_${{ matrix.config.framework }}_${{ matrix.config.report }}_test_reports + path: reports + + diff --git a/tests/modular_pipelines/stable_diffusion_xl/__init__.py b/tests/modular_pipelines/stable_diffusion_xl/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py b/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py index b8acca687c..b8a9a0c9a6 100644 --- a/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py +++ b/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py @@ -101,7 +101,7 @@ class SDXLModularTests: assert ( np.abs(image_slice.flatten() - expected_slice).max() < expected_max_diff - ), f"Image Slice does not match expected slice" + ), "Image Slice does not match expected slice" class SDXLModularIPAdapterTests: @@ -114,30 +114,20 @@ class SDXLModularIPAdapterTests: parameters = blocks.input_names assert issubclass(self.pipeline_class, ModularIPAdapterMixin) - self.assertIn( - "ip_adapter_image", - parameters, - "`ip_adapter_image` argument must be supported by the `__call__` method", - ) - self.assertIn( - "ip_adapter", - blocks.sub_blocks, - "pipeline must contain an IPAdapter block", - ) + assert ( + "ip_adapter_image" in parameters + ), "`ip_adapter_image` argument must be supported by the `__call__` method" + assert "ip_adapter" in blocks.sub_blocks, "pipeline must contain an IPAdapter block" _ = blocks.sub_blocks.pop("ip_adapter") parameters = blocks.input_names intermediate_parameters = blocks.intermediate_input_names - self.assertNotIn( - "ip_adapter_image", - parameters, - "`ip_adapter_image` argument must be removed from the `__call__` method", - ) - self.assertNotIn( - "ip_adapter_image_embeds", - intermediate_parameters, - "`ip_adapter_image_embeds` argument must be supported by the `__call__` method", - ) + assert ( + "ip_adapter_image" not in parameters + ), "`ip_adapter_image` argument must be removed from the `__call__` method" + assert ( + "ip_adapter_image_embeds" not in intermediate_parameters + ), "`ip_adapter_image_embeds` argument must be supported by the `__call__` method" def _get_dummy_image_embeds(self, cross_attention_dim: int = 32): return torch.randn((1, 1, cross_attention_dim), device=torch_device) @@ -213,14 +203,10 @@ class SDXLModularIPAdapterTests: max_diff_without_adapter_scale = np.abs(output_without_adapter_scale - output_without_adapter).max() max_diff_with_adapter_scale = np.abs(output_with_adapter_scale - output_without_adapter).max() - self.assertLess( - max_diff_without_adapter_scale, - expected_max_diff, - "Output without ip-adapter must be same as normal inference", - ) - self.assertGreater( - max_diff_with_adapter_scale, 1e-2, "Output with ip-adapter must be different from normal inference" - ) + assert ( + max_diff_without_adapter_scale < expected_max_diff + ), "Output without ip-adapter must be same as normal inference" + assert max_diff_with_adapter_scale > 1e-2, "Output with ip-adapter must be different from normal inference" # 2. Multi IP-Adapter test cases adapter_state_dict_1 = create_ip_adapter_state_dict(pipe.unet) @@ -249,16 +235,12 @@ class SDXLModularIPAdapterTests: output_without_multi_adapter_scale - output_without_adapter ).max() max_diff_with_multi_adapter_scale = np.abs(output_with_multi_adapter_scale - output_without_adapter).max() - self.assertLess( - max_diff_without_multi_adapter_scale, - expected_max_diff, - "Output without multi-ip-adapter must be same as normal inference", - ) - self.assertGreater( - max_diff_with_multi_adapter_scale, - 1e-2, - "Output with multi-ip-adapter scale must be different from normal inference", - ) + assert ( + max_diff_without_multi_adapter_scale < expected_max_diff + ), "Output without multi-ip-adapter must be same as normal inference" + assert ( + max_diff_with_multi_adapter_scale > 1e-2 + ), "Output with multi-ip-adapter scale must be different from normal inference" class SDXLModularControlNetTests: @@ -270,16 +252,10 @@ class SDXLModularControlNetTests: blocks = self.pipeline_blocks_class() parameters = blocks.input_names - self.assertIn( - "control_image", - parameters, - "`control_image` argument must be supported by the `__call__` method", - ) - self.assertIn( - "controlnet_conditioning_scale", - parameters, - "`controlnet_conditioning_scale` argument must be supported by the `__call__` method", - ) + assert "control_image" in parameters, "`control_image` argument must be supported by the `__call__` method" + assert ( + "controlnet_conditioning_scale" in parameters + ), "`controlnet_conditioning_scale` argument must be supported by the `__call__` method" def _modify_inputs_for_controlnet_test(self, inputs: Dict[str, Any]): controlnet_embedder_scale_factor = 2 @@ -325,14 +301,10 @@ class SDXLModularControlNetTests: max_diff_without_controlnet_scale = np.abs(output_without_controlnet_scale - output_without_controlnet).max() max_diff_with_controlnet_scale = np.abs(output_with_controlnet_scale - output_without_controlnet).max() - self.assertLess( - max_diff_without_controlnet_scale, - expected_max_diff, - "Output without controlnet must be same as normal inference", - ) - self.assertGreater( - max_diff_with_controlnet_scale, 1e-2, "Output with controlnet must be different from normal inference" - ) + assert ( + max_diff_without_controlnet_scale < expected_max_diff + ), "Output without controlnet must be same as normal inference" + assert max_diff_with_controlnet_scale > 1e-2, "Output with controlnet must be different from normal inference" def test_controlnet_cfg(self): pipe = self.get_pipeline() @@ -354,7 +326,7 @@ class SDXLModularControlNetTests: assert out_cfg.shape == out_no_cfg.shape max_diff = np.abs(out_cfg - out_no_cfg).max() - self.assertGreater(max_diff, 1e-2) + assert max_diff > 1e-2, "Output with CFG must be different from normal inference" class SDXLModularGuiderTests: @@ -378,7 +350,7 @@ class SDXLModularGuiderTests: assert out_cfg.shape == out_no_cfg.shape max_diff = np.abs(out_cfg - out_no_cfg).max() - self.assertGreater(max_diff, 1e-2) + assert max_diff > 1e-2, "Output with CFG must be different from normal inference" class SDXLModularPipelineFastTests( diff --git a/tests/modular_pipelines/test_modular_pipelines_common.py b/tests/modular_pipelines/test_modular_pipelines_common.py index e150a1e59e..1bb9bcf2cb 100644 --- a/tests/modular_pipelines/test_modular_pipelines_common.py +++ b/tests/modular_pipelines/test_modular_pipelines_common.py @@ -2,12 +2,12 @@ import gc import unittest from typing import Callable, Union -from diffusers.utils.dummy_pt_objects import ModularPipeline, ModularPipelineBlocks import numpy as np import torch import diffusers from diffusers.utils import logging +from diffusers.utils.dummy_pt_objects import ModularPipeline, ModularPipelineBlocks from diffusers.utils.testing_utils import ( backend_empty_cache, numpy_cosine_similarity_distance, @@ -142,7 +142,7 @@ class ModularPipelineTesterMixin: optional_parameters = pipe.default_call_parameters def _check_for_parameters(parameters, expected_parameters, param_type): - remaining_parameters = set(param for param in parameters if param not in expected_parameters) + remaining_parameters = {param for param in parameters if param not in expected_parameters} assert ( len(remaining_parameters) == 0 ), f"Required {param_type} parameters not present: {remaining_parameters}" @@ -188,7 +188,7 @@ class ModularPipelineTesterMixin: output = pipe(**batched_input, output="images") assert len(output) == batch_size, "Output is different from expected batch size" - def test_batch_inference_identical_to_single( + def test_inference_batch_single_identical( self, batch_size=2, expected_max_diff=1e-4, @@ -283,16 +283,16 @@ class ModularPipelineTesterMixin: pipe.set_progress_bar_config(disable=None) pipe.to("cpu") - output = pipe(**self.get_dummy_inputs("cpu"), output="np") + output = pipe(**self.get_dummy_inputs("cpu"), output="images") assert np.isnan(to_np(output)).sum() == 0, "CPU Inference returns NaN" @require_accelerator - def test_inferece_is_not_nan(self): + def test_inference_is_not_nan(self): pipe = self.get_pipeline() pipe.set_progress_bar_config(disable=None) pipe.to(torch_device) - output = pipe(**self.get_dummy_inputs(torch_device), output="np") + output = pipe(**self.get_dummy_inputs(torch_device), output="images") assert np.isnan(to_np(output)).sum() == 0, "Accelerator Inference returns NaN" def test_num_images_per_prompt(self): diff --git a/tests/pipelines/test_modular_pipelines_common.py b/tests/pipelines/test_modular_pipelines_common.py deleted file mode 100644 index 4bd45e207b..0000000000 --- a/tests/pipelines/test_modular_pipelines_common.py +++ /dev/null @@ -1,369 +0,0 @@ -import gc -import unittest -from typing import Callable, Union - -import numpy as np -import torch - -import diffusers -from diffusers import ( - DiffusionPipeline, -) -from diffusers.utils import logging -from diffusers.utils.testing_utils import ( - backend_empty_cache, - numpy_cosine_similarity_distance, - require_accelerator, - require_torch, - torch_device, -) - - -def to_np(tensor): - if isinstance(tensor, torch.Tensor): - tensor = tensor.detach().cpu().numpy() - - return tensor - - -@require_torch -class ModularPipelineTesterMixin: - """ - This mixin is designed to be used with unittest.TestCase classes. - It provides a set of common tests for each modular pipeline, - including: - - test_pipeline_call_signature: check if the pipeline's __call__ method has all required parameters - - test_inference_batch_consistent: check if the pipeline's __call__ method can handle batch inputs - - test_inference_batch_single_identical: check if the pipeline's __call__ method can handle single input - - test_float16_inference: check if the pipeline's __call__ method can handle float16 inputs - - test_to_device: check if the pipeline's __call__ method can handle different devices - """ - - # Canonical parameters that are passed to `__call__` regardless - # of the type of pipeline. They are always optional and have common - # sense default values. - required_optional_params = frozenset( - [ - "num_inference_steps", - "num_images_per_prompt", - "latents", - "output_type", - ] - ) - # this is modular specific: generator needs to be a intermediate input because it's mutable - required_intermediate_params = frozenset( - [ - "generator", - ] - ) - - def get_generator(self, seed): - device = torch_device if torch_device != "mps" else "cpu" - generator = torch.Generator(device).manual_seed(seed) - return generator - - @property - def pipeline_class(self) -> Union[Callable, DiffusionPipeline]: - raise NotImplementedError( - "You need to set the attribute `pipeline_class = ClassNameOfPipeline` in the child test class. " - "See existing pipeline tests for reference." - ) - - @property - def repo(self) -> str: - raise NotImplementedError( - "You need to set the attribute `repo` in the child test class. See existing pipeline tests for reference." - ) - - @property - def pipeline_blocks_class(self) -> Union[Callable, DiffusionPipeline]: - raise NotImplementedError( - "You need to set the attribute `pipeline_blocks_class = ClassNameOfPipelineBlocks` in the child test class. " - "See existing pipeline tests for reference." - ) - - def get_pipeline(self): - raise NotImplementedError( - "You need to implement `get_pipeline(self)` in the child test class. " - "See existing pipeline tests for reference." - ) - - def get_dummy_inputs(self, device, seed=0): - raise NotImplementedError( - "You need to implement `get_dummy_inputs(self, device, seed)` in the child test class. " - "See existing pipeline tests for reference." - ) - - @property - def params(self) -> frozenset: - raise NotImplementedError( - "You need to set the attribute `params` in the child test class. " - "`params` are checked for if all values are present in `__call__`'s signature." - " You can set `params` using one of the common set of parameters defined in `pipeline_params.py`" - " e.g., `TEXT_TO_IMAGE_PARAMS` defines the common parameters used in text to " - "image pipelines, including prompts and prompt embedding overrides." - "If your pipeline's set of arguments has minor changes from one of the common sets of arguments, " - "do not make modifications to the existing common sets of arguments. I.e. a text to image pipeline " - "with non-configurable height and width arguments should set the attribute as " - "`params = TEXT_TO_IMAGE_PARAMS - {'height', 'width'}`. " - "See existing pipeline tests for reference." - ) - - @property - def batch_params(self) -> frozenset: - raise NotImplementedError( - "You need to set the attribute `batch_params` in the child test class. " - "`batch_params` are the parameters required to be batched when passed to the pipeline's " - "`__call__` method. `pipeline_params.py` provides some common sets of parameters such as " - "`TEXT_TO_IMAGE_BATCH_PARAMS`, `IMAGE_VARIATION_BATCH_PARAMS`, etc... If your pipeline's " - "set of batch arguments has minor changes from one of the common sets of batch arguments, " - "do not make modifications to the existing common sets of batch arguments. I.e. a text to " - "image pipeline `negative_prompt` is not batched should set the attribute as " - "`batch_params = TEXT_TO_IMAGE_BATCH_PARAMS - {'negative_prompt'}`. " - "See existing pipeline tests for reference." - ) - - def setUp(self): - # clean up the VRAM before each test - super().setUp() - torch.compiler.reset() - gc.collect() - backend_empty_cache(torch_device) - - def tearDown(self): - # clean up the VRAM after each test in case of CUDA runtime errors - super().tearDown() - torch.compiler.reset() - gc.collect() - backend_empty_cache(torch_device) - - def test_pipeline_call_signature(self): - pipe = self.get_pipeline() - parameters = pipe.blocks.input_names - optional_parameters = pipe.default_call_parameters - intermediate_parameters = pipe.blocks.intermediate_input_names - - remaining_required_parameters = set() - - for param in self.params: - if param not in parameters: - remaining_required_parameters.add(param) - - self.assertTrue( - len(remaining_required_parameters) == 0, - f"Required parameters not present: {remaining_required_parameters}", - ) - - remaining_required_intermediate_parameters = set() - - for param in self.required_intermediate_params: - if param not in intermediate_parameters: - remaining_required_intermediate_parameters.add(param) - - self.assertTrue( - len(remaining_required_intermediate_parameters) == 0, - f"Required intermediate parameters not present: {remaining_required_intermediate_parameters}", - ) - - remaining_required_optional_parameters = set() - - for param in self.required_optional_params: - if param not in optional_parameters: - remaining_required_optional_parameters.add(param) - - self.assertTrue( - len(remaining_required_optional_parameters) == 0, - f"Required optional parameters not present: {remaining_required_optional_parameters}", - ) - - def test_inference_batch_consistent(self, batch_sizes=[2]): - self._test_inference_batch_consistent(batch_sizes=batch_sizes) - - def _test_inference_batch_consistent( - self, batch_sizes=[2], additional_params_copy_to_batched_inputs=["num_inference_steps"], batch_generator=True - ): - pipe = self.get_pipeline() - pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(torch_device) - inputs["generator"] = self.get_generator(0) - - logger = logging.get_logger(pipe.__module__) - logger.setLevel(level=diffusers.logging.FATAL) - - # prepare batched inputs - batched_inputs = [] - for batch_size in batch_sizes: - batched_input = {} - batched_input.update(inputs) - - for name in self.batch_params: - if name not in inputs: - continue - - value = inputs[name] - if name == "prompt": - len_prompt = len(value) - # make unequal batch sizes - batched_input[name] = [value[: len_prompt // i] for i in range(1, batch_size + 1)] - - # make last batch super long - batched_input[name][-1] = 100 * "very long" - - else: - batched_input[name] = batch_size * [value] - - if batch_generator and "generator" in inputs: - batched_input["generator"] = [self.get_generator(i) for i in range(batch_size)] - - if "batch_size" in inputs: - batched_input["batch_size"] = batch_size - - batched_inputs.append(batched_input) - - logger.setLevel(level=diffusers.logging.WARNING) - for batch_size, batched_input in zip(batch_sizes, batched_inputs): - output = pipe(**batched_input, output="images") - assert len(output) == batch_size - - def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff=1e-4): - self._test_inference_batch_single_identical(batch_size=batch_size, expected_max_diff=expected_max_diff) - - def _test_inference_batch_single_identical( - self, - batch_size=2, - expected_max_diff=1e-4, - additional_params_copy_to_batched_inputs=["num_inference_steps"], - ): - pipe = self.get_pipeline() - for component in pipe.components.values(): - if hasattr(component, "set_default_attn_processor"): - component.set_default_attn_processor() - - pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - inputs = self.get_dummy_inputs(torch_device) - # Reset generator in case it is has been used in self.get_dummy_inputs - inputs["generator"] = self.get_generator(0) - - logger = logging.get_logger(pipe.__module__) - logger.setLevel(level=diffusers.logging.FATAL) - - # batchify inputs - batched_inputs = {} - batched_inputs.update(inputs) - - for name in self.batch_params: - if name not in inputs: - continue - - value = inputs[name] - if name == "prompt": - len_prompt = len(value) - batched_inputs[name] = [value[: len_prompt // i] for i in range(1, batch_size + 1)] - batched_inputs[name][-1] = 100 * "very long" - - else: - batched_inputs[name] = batch_size * [value] - - if "generator" in inputs: - batched_inputs["generator"] = [self.get_generator(i) for i in range(batch_size)] - - if "batch_size" in inputs: - batched_inputs["batch_size"] = batch_size - - for arg in additional_params_copy_to_batched_inputs: - batched_inputs[arg] = inputs[arg] - - output = pipe(**inputs, output="images") - output_batch = pipe(**batched_inputs, output="images") - - assert output_batch.shape[0] == batch_size - - max_diff = np.abs(to_np(output_batch[0]) - to_np(output[0])).max() - assert max_diff < expected_max_diff - - @unittest.skipIf(torch_device not in ["cuda", "xpu"], reason="float16 requires CUDA or XPU") - @require_accelerator - def test_float16_inference(self, expected_max_diff=5e-2): - pipe = self.get_pipeline(torch_dtype=torch.float32) - for component in pipe.components.values(): - if hasattr(component, "set_default_attn_processor"): - component.set_default_attn_processor() - - pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - pipe_fp16 = self.get_pipeline(torch_dtype=torch.float16) - for component in pipe_fp16.components.values(): - if hasattr(component, "set_default_attn_processor"): - component.set_default_attn_processor() - pipe_fp16.to(torch_device, torch.float16) - pipe_fp16.set_progress_bar_config(disable=None) - - inputs = self.get_dummy_inputs(torch_device) - # Reset generator in case it is used inside dummy inputs - if "generator" in inputs: - inputs["generator"] = self.get_generator(0) - output = pipe(**inputs, output="images") - - fp16_inputs = self.get_dummy_inputs(torch_device) - # Reset generator in case it is used inside dummy inputs - if "generator" in fp16_inputs: - fp16_inputs["generator"] = self.get_generator(0) - output_fp16 = pipe_fp16(**fp16_inputs, output="images") - - if isinstance(output, torch.Tensor): - output = output.cpu() - output_fp16 = output_fp16.cpu() - - max_diff = numpy_cosine_similarity_distance(output.flatten(), output_fp16.flatten()) - assert max_diff < expected_max_diff - - @require_accelerator - def test_to_device(self): - pipe = self.get_pipeline() - pipe.set_progress_bar_config(disable=None) - - pipe.to("cpu") - model_devices = [ - component.device.type for component in pipe.components.values() if hasattr(component, "device") - ] - self.assertTrue(all(device == "cpu" for device in model_devices)) - - output_cpu = pipe(**self.get_dummy_inputs("cpu"), output="images") - self.assertTrue(np.isnan(output_cpu).sum() == 0) - - pipe.to(torch_device) - model_devices = [ - component.device.type for component in pipe.components.values() if hasattr(component, "device") - ] - self.assertTrue(all(device == torch_device for device in model_devices)) - - output_device = pipe(**self.get_dummy_inputs(torch_device), output="images") - self.assertTrue(np.isnan(to_np(output_device)).sum() == 0) - - def test_num_images_per_prompt(self): - pipe = self.get_pipeline() - - if "num_images_per_prompt" not in pipe.blocks.input_names: - return - - pipe = pipe.to(torch_device) - pipe.set_progress_bar_config(disable=None) - - batch_sizes = [1, 2] - num_images_per_prompts = [1, 2] - - for batch_size in batch_sizes: - for num_images_per_prompt in num_images_per_prompts: - inputs = self.get_dummy_inputs(torch_device) - - for key in inputs.keys(): - if key in self.batch_params: - inputs[key] = batch_size * [inputs[key]] - - images = pipe(**inputs, num_images_per_prompt=num_images_per_prompt, output="images") - - assert images.shape[0] == batch_size * num_images_per_prompt