From 0eb5dce09437be73618cbb7c3d712727d938bae6 Mon Sep 17 00:00:00 2001 From: IlyasMoutawwakil Date: Thu, 25 Jul 2024 16:33:01 +0200 Subject: [PATCH] fix clip text model with projection not outputting hidden states --- optimum/exporters/utils.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/optimum/exporters/utils.py b/optimum/exporters/utils.py index 51ef8da0af8..e2125736c4d 100644 --- a/optimum/exporters/utils.py +++ b/optimum/exporters/utils.py @@ -117,10 +117,11 @@ def _get_submodels_for_export_diffusion( models_for_export = {} # Text encoder - if pipeline.text_encoder is not None: + text_encoder = getattr(pipeline, "text_encoder", None) + if text_encoder is not None: if is_stable_diffusion_xl: - pipeline.text_encoder.config.output_hidden_states = True - models_for_export["text_encoder"] = pipeline.text_encoder + text_encoder.config.output_hidden_states = True + models_for_export["text_encoder"] = text_encoder # U-NET # ONNX export of torch.nn.functional.scaled_dot_product_attention not supported for < v2.1.0 @@ -150,6 +151,8 @@ def _get_submodels_for_export_diffusion( text_encoder_2 = getattr(pipeline, "text_encoder_2", None) if text_encoder_2 is not None: + text_encoder_2.config.output_hidden_states = True + text_encoder_2.text_model.config.output_hidden_states = True models_for_export["text_encoder_2"] = text_encoder_2 return models_for_export