fixes

huggingface · Oct 23, 2024 · aa74f63 · aa74f63
1 parent 4518691
commit aa74f63
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 5 deletions.
diff --git a/optimum/exporters/onnx/model_configs.py b/optimum/exporters/onnx/model_configs.py
@@ -1106,7 +1106,7 @@ def ordered_inputs(self, model) -> Dict[str, Dict[int, str]]:
 
 
 class VaeEncoderOnnxConfig(VisionOnnxConfig):
-    ATOL_FOR_VALIDATION = 1e-4
+    ATOL_FOR_VALIDATION = 3e-4
     # The ONNX export of a CLIPText architecture, an other Stable Diffusion component, needs the Trilu
     # operator support, available since opset 14
     DEFAULT_ONNX_OPSET = 14

diff --git a/optimum/exporters/tasks.py b/optimum/exporters/tasks.py
@@ -1182,15 +1182,17 @@ class TasksManager:
         "transformers": _SUPPORTED_MODEL_TYPE,
     }
     _UNSUPPORTED_CLI_MODEL_TYPE = {
-        # diffusers submodels
+        # diffusers model types
         "clip-text-model",
         "clip-text-model-with-projection",
-        "flux-transformer-model",
+        "flux-transformer-2d-model",
         "sd3-transformer-2d-model",
         "t5-encoder-model",
         "unet-2d-condition",
         "vae-encoder",
         "vae-decoder",
+        # redundant model types
+        "trocr",  # same as vision-encoder-decoder
     }
     _SUPPORTED_CLI_MODEL_TYPE = (
         set(_SUPPORTED_MODEL_TYPE.keys())

diff --git a/tests/exporters/onnx/test_onnx_export.py b/tests/exporters/onnx/test_onnx_export.py
@@ -299,15 +299,13 @@ def _onnx_export_diffusion_models(self, model_type: str, model_name: str, device
         with TemporaryDirectory() as tmpdirname:
             _, onnx_outputs = export_models(
                 models_and_onnx_configs=models_and_onnx_configs,
-                opset=14,
                 output_dir=Path(tmpdirname),
                 device=device,
             )
             validate_models_outputs(
                 models_and_onnx_configs=models_and_onnx_configs,
                 onnx_named_outputs=onnx_outputs,
                 output_dir=Path(tmpdirname),
-                atol=1e-4,
                 use_subprocess=False,
             )