Code cleanup-2

Signed-off-by: Amit Raj <[email protected]>

Author: quic-amitraj

QEfficient/diffusers/pipelines/pipeline_utils.py

@@ -38,12 +38,16 @@ def get_onnx_config(self):
 
         example_inputs = {
             "input_ids": torch.zeros((bs, seq_len), dtype=torch.int64),
-            "output_hidden_states": True,
         }
 
         dynamic_axes = {"input_ids": {0: "batch_size", 1: "seq_len"}}
 
         output_names = ["pooler_output", "last_hidden_state"]
+        if self.model.__class__.__name__ == "T5EncoderModel":
+            output_names = ["last_hidden_state"]
+        else:
+            example_inputs["output_hidden_states"] = (True,)
+
         return example_inputs, dynamic_axes, output_names
 
     def export(self, inputs, output_names, dynamic_axes, export_dir=None):

QEfficient/diffusers/pipelines/stable_diffusion_3/pipeline_stable_diffusion3.py

@@ -242,6 +242,7 @@ def compile(
         )
 
         self.text_encoder_compile_path = self.text_encoder._compile(
+            onnx_path,
             compile_dir,
             compile_only=True,
             specializations=specializations_text_encoder,