diff --git a/optimum/exporters/onnx/__init__.py b/optimum/exporters/onnx/__init__.py index 609096e37ef..6b99e484578 100644 --- a/optimum/exporters/onnx/__init__.py +++ b/optimum/exporters/onnx/__init__.py @@ -31,7 +31,7 @@ "utils": [ "get_decoder_models_for_export", "get_encoder_decoder_models_for_export", - "get_stable_diffusion_models_for_export", + "get_diffusion_models_for_export", "MODEL_TYPES_REQUIRING_POSITION_IDS", ], "__main__": ["main_export"], @@ -50,7 +50,7 @@ from .utils import ( get_decoder_models_for_export, get_encoder_decoder_models_for_export, - get_stable_diffusion_models_for_export, + get_diffusion_models_for_export, MODEL_TYPES_REQUIRING_POSITION_IDS, ) from .__main__ import main_export diff --git a/optimum/exporters/utils.py b/optimum/exporters/utils.py index 881741ef3dc..b0f20d330e1 100644 --- a/optimum/exporters/utils.py +++ b/optimum/exporters/utils.py @@ -87,7 +87,7 @@ def _get_submodels_for_export_diffusion( - pipeline: "DiffusionPipeline" + pipeline: "DiffusionPipeline", ) -> Dict[str, Union["PreTrainedModel", "ModelMixin"]]: """ Returns the components of a Stable Diffusion model. diff --git a/tests/exporters/exporters_utils.py b/tests/exporters/exporters_utils.py index 0c52754ff60..254cd40194a 100644 --- a/tests/exporters/exporters_utils.py +++ b/tests/exporters/exporters_utils.py @@ -294,7 +294,7 @@ "roberta": "roberta-base", } -PYTORCH_STABLE_DIFFUSION_MODEL = { +PYTORCH_DIFFUSION_MODEL = { "stable-diffusion": "hf-internal-testing/tiny-stable-diffusion-torch", "stable-diffusion-xl": "echarlaix/tiny-random-stable-diffusion-xl", } diff --git a/tests/exporters/onnx/test_exporters_onnx_cli.py b/tests/exporters/onnx/test_exporters_onnx_cli.py index 667206b0068..ed611ade04e 100644 --- a/tests/exporters/onnx/test_exporters_onnx_cli.py +++ b/tests/exporters/onnx/test_exporters_onnx_cli.py @@ -41,9 +41,9 @@ from ..exporters_utils import ( NO_DYNAMIC_AXES_EXPORT_SHAPES_TRANSFORMERS, + PYTORCH_DIFFUSION_MODEL, PYTORCH_EXPORT_MODELS_TINY, PYTORCH_SENTENCE_TRANSFORMERS_MODEL, - PYTORCH_STABLE_DIFFUSION_MODEL, PYTORCH_TIMM_MODEL, PYTORCH_TIMM_MODEL_NO_DYNAMIC_AXES, PYTORCH_TRANSFORMERS_MODEL_NO_DYNAMIC_AXES, @@ -252,29 +252,29 @@ def _onnx_export_no_dynamic_axes( except MinimumVersionError as e: pytest.skip(f"Skipping due to minimum version requirements not met. Full error: {e}") - @parameterized.expand(PYTORCH_STABLE_DIFFUSION_MODEL.items()) + @parameterized.expand(PYTORCH_DIFFUSION_MODEL.items()) @require_torch @require_vision @require_diffusers - def test_exporters_cli_pytorch_cpu_stable_diffusion(self, model_type: str, model_name: str): + def test_exporters_cli_pytorch_cpu_diffusion(self, model_type: str, model_name: str): self._onnx_export(model_name, model_type) - @parameterized.expand(PYTORCH_STABLE_DIFFUSION_MODEL.items()) + @parameterized.expand(PYTORCH_DIFFUSION_MODEL.items()) @require_torch_gpu @require_vision @require_diffusers @slow @pytest.mark.run_slow - def test_exporters_cli_pytorch_gpu_stable_diffusion(self, model_type: str, model_name: str): + def test_exporters_cli_pytorch_gpu_diffusion(self, model_type: str, model_name: str): self._onnx_export(model_name, model_type, device="cuda") - @parameterized.expand(PYTORCH_STABLE_DIFFUSION_MODEL.items()) + @parameterized.expand(PYTORCH_DIFFUSION_MODEL.items()) @require_torch_gpu @require_vision @require_diffusers @slow @pytest.mark.run_slow - def test_exporters_cli_fp16_stable_diffusion(self, model_type: str, model_name: str): + def test_exporters_cli_fp16_diffusion(self, model_type: str, model_name: str): self._onnx_export(model_name, model_type, device="cuda", fp16=True) @parameterized.expand( @@ -594,7 +594,7 @@ def test_trust_remote_code(self): check=True, ) - def test_stable_diffusion(self): + def test_diffusion(self): with TemporaryDirectory() as tmpdirname: subprocess.run( f"python3 -m optimum.exporters.onnx --model hf-internal-testing/tiny-stable-diffusion-torch --task stable-diffusion {tmpdirname}", diff --git a/tests/exporters/onnx/test_onnx_export.py b/tests/exporters/onnx/test_onnx_export.py index 9eddc4c86de..d1471aa218a 100644 --- a/tests/exporters/onnx/test_onnx_export.py +++ b/tests/exporters/onnx/test_onnx_export.py @@ -32,8 +32,8 @@ OnnxConfigWithPast, export_models, get_decoder_models_for_export, + get_diffusion_models_for_export, get_encoder_decoder_models_for_export, - get_stable_diffusion_models_for_export, main_export, onnx_export_from_model, validate_models_outputs, @@ -48,9 +48,9 @@ from optimum.utils.testing_utils import grid_parameters, require_diffusers from ..exporters_utils import ( + PYTORCH_DIFFUSION_MODEL, PYTORCH_EXPORT_MODELS_TINY, PYTORCH_SENTENCE_TRANSFORMERS_MODEL, - PYTORCH_STABLE_DIFFUSION_MODEL, PYTORCH_TIMM_MODEL, TENSORFLOW_EXPORT_MODELS, VALIDATE_EXPORT_ON_SHAPES_SLOW, @@ -294,7 +294,7 @@ def _onnx_export( def _onnx_export_sd(self, model_type: str, model_name: str, device="cpu"): pipeline = TasksManager.get_model_from_task(model_type, model_name, device=device) - models_and_onnx_configs = get_stable_diffusion_models_for_export(pipeline) + models_and_onnx_configs = get_diffusion_models_for_export(pipeline) output_names = [os.path.join(name_dir, ONNX_WEIGHTS_NAME) for name_dir in models_and_onnx_configs] model, _ = models_and_onnx_configs["vae_encoder"] model.forward = lambda sample: {"latent_sample": model.encode(x=sample)["latent_dist"].parameters} @@ -398,14 +398,14 @@ def test_tensorflow_export( self._onnx_export(test_name, model_type, model_name, task, onnx_config_class_constructor, monolith=monolith) - @parameterized.expand(PYTORCH_STABLE_DIFFUSION_MODEL.items()) + @parameterized.expand(PYTORCH_DIFFUSION_MODEL.items()) @require_torch @require_vision @require_diffusers - def test_pytorch_export_for_stable_diffusion_models(self, model_type, model_name): + def test_pytorch_export_for_diffusion_models(self, model_type, model_name): self._onnx_export_sd(model_type, model_name) - @parameterized.expand(PYTORCH_STABLE_DIFFUSION_MODEL.items()) + @parameterized.expand(PYTORCH_DIFFUSION_MODEL.items()) @require_torch @require_vision @require_diffusers @@ -413,7 +413,7 @@ def test_pytorch_export_for_stable_diffusion_models(self, model_type, model_name @slow @pytest.mark.run_slow @pytest.mark.gpu_test - def test_pytorch_export_for_stable_diffusion_models_cuda(self, model_type, model_name): + def test_pytorch_export_for_diffusion_models_cuda(self, model_type, model_name): self._onnx_export_sd(model_type, model_name, device="cuda")