fix

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

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.

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",
 }

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}",

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,22 +398,22 @@ 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
     @require_torch_gpu
     @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")