Skip to content

Commit

Permalink
fix clip text model with projection not outputting hidden states
Browse files Browse the repository at this point in the history
  • Loading branch information
IlyasMoutawwakil committed Jul 25, 2024
1 parent c761026 commit 0eb5dce
Showing 1 changed file with 6 additions and 3 deletions.
9 changes: 6 additions & 3 deletions optimum/exporters/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 0eb5dce

Please sign in to comment.