Replaced output dict with dataclass to make it more user friendly

Signed-off-by: Amit Raj <amitraj@qti.qualcommm.com>

Author: Amit Raj

QEfficient/diffusers/pipelines/flux/pipeline_flux.py

@@ -21,7 +21,12 @@
     QEffTextEncoder,
     QEffVAE,
 )
-from QEfficient.diffusers.pipelines.pipeline_utils import QEffPipelineOutput, config_manager, set_module_device_ids
+from QEfficient.diffusers.pipelines.pipeline_utils import (
+    ModulePerf,
+    QEffPipelineOutput,
+    config_manager,
+    set_module_device_ids,
+)
 from QEfficient.generation.cloud_infer import QAICInferenceSession
 
 
@@ -42,13 +47,13 @@ def __init__(self, model, use_onnx_function, *args, **kwargs):
         self.vae_decode = QEffVAE(model, "decoder")
         self.use_onnx_function = use_onnx_function
 
-        # Add all modules of FluxPipeline
-        self.has_module = [
-            ("text_encoder", self.text_encoder),
-            ("text_encoder_2", self.text_encoder_2),
-            ("transformer", self.transformer),
-            ("vae_decoder", self.vae_decode),
-        ]
+        # All modules of FluxPipeline stored in a dictionary for easy access and iteration
+        self.modules = {
+            "text_encoder": self.text_encoder,
+            "text_encoder_2": self.text_encoder_2,
+            "transformer": self.transformer,
+            "vae_decoder": self.vae_decode,
+        }
 
         self.tokenizer = model.tokenizer
         self.text_encoder.tokenizer = model.tokenizer
@@ -127,7 +132,7 @@ def export(self, export_dir: Optional[str] = None) -> str:
             :str: Path of the generated ``ONNX`` graph.
         """
 
-        for module_name, module_obj in self.has_module:
+        for module_name, module_obj in self.modules.items():
             example_inputs_text_encoder, dynamic_axes_text_encoder, output_names_text_encoder = (
                 module_obj.get_onnx_config()
             )
@@ -183,7 +188,7 @@ def compile(
         if self.custom_config is None:
             config_manager(self, config_source=compile_config)
 
-        for module_name, module_obj in self.has_module:
+        for module_name, module_obj in self.modules.items():
             # Get specialization values directly from config
             module_config = self.custom_config["modules"]
             specializations = [module_config[module_name]["specializations"]]
@@ -256,19 +261,18 @@ def _get_t5_prompt_embeds(
         self.text_encoder_2.qpc_session.set_buffers(text_encoder_2_output)
 
         aic_text_input = {"input_ids": text_input_ids.numpy().astype(np.int64)}
-        import time
 
         start_t5_time = time.time()
         prompt_embeds = torch.tensor(self.text_encoder_2.qpc_session.run(aic_text_input)["last_hidden_state"])
         end_t5_time = time.time()
-        self.text_encoder_2.inference_time = end_t5_time - start_t5_time
+        text_encoder_2_perf = end_t5_time - start_t5_time
 
         _, seq_len, _ = prompt_embeds.shape
         # duplicate text embeddings and attention mask for each generation per prompt, using mps friendly method
         prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
         prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
 
-        return prompt_embeds
+        return prompt_embeds, text_encoder_2_perf
 
     def _get_clip_prompt_embeds(
         self,
@@ -322,20 +326,17 @@ def _get_clip_prompt_embeds(
 
         aic_text_input = {"input_ids": text_input_ids.numpy().astype(np.int64)}
 
-        import time
-
-        global start_text_encoder_time
         start_text_encoder_time = time.time()
         aic_embeddings = self.text_encoder.qpc_session.run(aic_text_input)
         end_text_encoder_time = time.time()
-        self.text_encoder.inference_time = end_text_encoder_time - start_text_encoder_time
+        text_encoder_perf = end_text_encoder_time - start_text_encoder_time
         prompt_embeds = torch.tensor(aic_embeddings["pooler_output"])
 
         # duplicate text embeddings for each generation per prompt, using mps friendly method
         prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt)
         prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1)
 
-        return prompt_embeds
+        return prompt_embeds, text_encoder_perf
 
     def encode_prompt(
         self,
@@ -378,20 +379,20 @@ def encode_prompt(
             prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
 
             # We only use the pooled prompt output from the CLIPTextModel
-            pooled_prompt_embeds = self._get_clip_prompt_embeds(
+            pooled_prompt_embeds, text_encoder_perf = self._get_clip_prompt_embeds(
                 prompt=prompt,
                 device_ids=self.text_encoder.device_ids,
                 num_images_per_prompt=num_images_per_prompt,
             )
-            prompt_embeds = self._get_t5_prompt_embeds(
+            prompt_embeds, text_encoder_2_perf = self._get_t5_prompt_embeds(
                 prompt=prompt_2,
                 num_images_per_prompt=num_images_per_prompt,
                 max_sequence_length=max_sequence_length,
                 device_ids=self.text_encoder_2.device_ids,
             )
 
         text_ids = torch.zeros(prompt_embeds.shape[1], 3)
-        return prompt_embeds, pooled_prompt_embeds, text_ids
+        return prompt_embeds, pooled_prompt_embeds, text_ids, [text_encoder_perf, text_encoder_2_perf]
 
     def __call__(
         self,
@@ -539,18 +540,15 @@ def __call__(
             negative_prompt_embeds is not None and negative_pooled_prompt_embeds is not None
         )
         do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
-        (
-            prompt_embeds,
-            pooled_prompt_embeds,
-            text_ids,
-        ) = self.encode_prompt(
+        (prompt_embeds, pooled_prompt_embeds, text_ids, text_encoder_perf) = self.encode_prompt(
             prompt=prompt,
             prompt_2=prompt_2,
             prompt_embeds=prompt_embeds,
             pooled_prompt_embeds=pooled_prompt_embeds,
             num_images_per_prompt=num_images_per_prompt,
             max_sequence_length=max_sequence_length,
         )
+
         if do_true_cfg:
             (
                 negative_prompt_embeds,
@@ -595,7 +593,7 @@ def __call__(
         }
 
         self.transformer.qpc_session.set_buffers(output_buffer)
-        self.transformer.inference_time = []
+        transformer_perf = []
         self.scheduler.set_begin_index(0)
         with self.progress_bar(total=num_inference_steps) as progress_bar:
             for i, t in enumerate(timesteps):
@@ -653,7 +651,7 @@ def __call__(
                 start_transformer_step_time = time.time()
                 outputs = self.transformer.qpc_session.run(inputs_aic)
                 end_transfromer_step_time = time.time()
-                self.transformer.inference_time.append(end_transfromer_step_time - start_transformer_step_time)
+                transformer_perf.append(end_transfromer_step_time - start_transformer_step_time)
 
                 noise_pred = torch.from_numpy(outputs["output"])
 
@@ -678,7 +676,6 @@ def __call__(
                 # call the callback, if provided
                 if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
                     progress_bar.update()
-
         if output_type == "latent":
             image = latents
         else:
@@ -704,14 +701,22 @@ def __call__(
             start_decode_time = time.time()
             image = self.vae_decode.qpc_session.run(inputs)
             end_decode_time = time.time()
-            self.vae_decode.inference_time = end_decode_time - start_decode_time
+            vae_decode_perf = end_decode_time - start_decode_time
             image_tensor = torch.from_numpy(image["sample"])
             image = self.image_processor.postprocess(image_tensor, output_type=output_type)
 
-            total_time_taken = end_decode_time - start_text_encoder_time
+            # Collect performance data in a dict
+            perf_data = {
+                "text_encoder": text_encoder_perf[0],
+                "text_encoder_2": text_encoder_perf[1],
+                "transformer": transformer_perf,
+                "vae_decoder": vae_decode_perf,
+            }
 
-        return QEffPipelineOutput(
-            pipeline=self,
-            images=image,
-            E2E_time=total_time_taken,
-        )
+            # Build performance metrics dynamically
+            perf_metrics = [ModulePerf(module_name=name, perf=perf_data[name]) for name in self.modules.keys()]
+
+            return QEffPipelineOutput(
+                pipeline_module=perf_metrics,
+                images=image,
+            )

QEfficient/diffusers/pipelines/pipeline_utils.py

@@ -7,16 +7,13 @@
 
 import os
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, List, Optional, Union
+from typing import List, Optional, Union
 
 import numpy as np
 import PIL.Image
 
 from QEfficient.utils._utils import load_json
 
-if TYPE_CHECKING:
-    from QEfficient.diffusers.pipelines.flux.pipeline_flux import QEFFFluxPipeline
-
 
 def config_manager(cls, config_source: Optional[str] = None):
     """
@@ -50,51 +47,58 @@ def set_module_device_ids(cls):
     from the configuration file to each module's device_ids attribute.
     """
     config_modules = cls.custom_config["modules"]
-    for module_name, module_obj in cls.has_module:
+    for module_name, module_obj in cls.modules.items():
         module_obj.device_ids = config_modules[module_name]["execute"]["device_ids"]
 
+@dataclass(frozen=True)
+class ModulePerf:
+    module_name: str
+    perf: int
+
 
-@dataclass
+@dataclass(frozen=True)
 class QEffPipelineOutput:
-    pipeline: "QEFFFluxPipeline"
+    pipeline_module: list[ModulePerf]
     images: Union[List[PIL.Image.Image], np.ndarray]
-    E2E_time: int
 
     def __repr__(self):
         output_str = "=" * 60 + "\n"
         output_str += "QEfficient Diffusers Pipeline Inference Report\n"
         output_str += "=" * 60 + "\n\n"
 
-        # End-to-End time
-        output_str += f"End-to-End Inference Time: {self.E2E_time:.4f} s\n\n"
-
         # Module-wise inference times
         output_str += "Module-wise Inference Times:\n"
         output_str += "-" * 60 + "\n"
 
-        # Iterate through all modules using has_module
-        for module_name, module_obj in self.pipeline.has_module:
-            if hasattr(module_obj, "inference_time"):
-                inference_time = module_obj.inference_time
-
-                # Format module name for display
-                display_name = module_name.replace("_", " ").title()
-
-                # Handle transformer specially as it has a list of times
-                if isinstance(inference_time, list) and len(inference_time) > 0:
-                    total_time = sum(inference_time)
-                    avg_time = total_time / len(inference_time)
-                    output_str += f"  {display_name:25s} {total_time:.4f} s\n"
-                    output_str += f"    - Total steps: {len(inference_time)}\n"
-                    output_str += f"    - Average per step:    {avg_time:.4f} s\n"
-                    output_str += f"    - Min step time:       {min(inference_time):.4f} s\n"
-                    output_str += f"    - Max step time:       {max(inference_time):.4f} s\n"
-                else:
-                    # Single inference time value
-                    output_str += f"  {display_name:25s} {inference_time:.4f} s\n"
+        # Calculate E2E time while iterating
+        e2e_time = 0
+        for module_perf in self.pipeline_module:
+            module_name = module_perf.module_name
+            inference_time = module_perf.perf
+
+            # Add to E2E time
+            e2e_time += sum(inference_time) if isinstance(inference_time, list) else inference_time
+
+            # Format module name for display
+            display_name = module_name.replace("_", " ").title()
+
+            # Handle transformer specially as it has a list of times
+            if isinstance(inference_time, list) and len(inference_time) > 0:
+                total_time = sum(inference_time)
+                avg_time = total_time / len(inference_time)
+                output_str += f"  {display_name:25s} {total_time:.4f} s\n"
+                output_str += f"    - Total steps: {len(inference_time)}\n"
+                output_str += f"    - Average per step:    {avg_time:.4f} s\n"
+                output_str += f"    - Min step time:       {min(inference_time):.4f} s\n"
+                output_str += f"    - Max step time:       {max(inference_time):.4f} s\n"
+            else:
+                # Single inference time value
+                output_str += f"  {display_name:25s} {inference_time:.4f} s\n"
 
         output_str += "-" * 60 + "\n\n"
 
+        # Print E2E time after all modules
+        output_str += f"End-to-End Inference Time: {e2e_time:.4f} s\n\n"
         output_str += "=" * 60 + "\n"
 
         return output_str

examples/diffusers/flux/flux_1_schnell.py

@@ -17,7 +17,7 @@
     max_sequence_length=256,
     generator=torch.manual_seed(42),
 )
-image = pipeline.images[0]
+image = output.images[0]
 image.save("cat_with_sign.png")
 
 print(output)

examples/diffusers/flux/flux_1_shnell_custom.py

@@ -112,7 +112,7 @@
     generator=torch.manual_seed(42),
 )
 
-image = output.images[0]
+images = output.images[0]
 # Save the generated image to disk
-image.save("girl_laughing.png")
+images.save("girl_laughing.png")
 print(output)