[Profile] Adding metrics for Diffusion/DiT Single diffusion Pipeline (#668)

Signed-off-by: Chen Yang <2082464740@qq.com>
Signed-off-by: erfgss <97771661+erfgss@users.noreply.github.com>

vllm_omni/diffusion/diffusion_engine.py

@@ -67,14 +67,19 @@ class DiffusionEngine:
             raise e
 
     def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
+        diffusion_engine_start_time = time.time()
         # Apply pre-processing if available
+        preprocess_time = 0.0
         if self.pre_process_func is not None:
             preprocess_start_time = time.time()
             request = self.pre_process_func(request)
             preprocess_time = time.time() - preprocess_start_time
             logger.info(f"Pre-processing completed in {preprocess_time:.4f} seconds")
 
+        exec_start_time = time.time()
         output = self.add_req_and_wait_for_response(request)
+        exec_total_time = time.time() - exec_start_time
+
         if output.error:
             raise Exception(f"{output.error}")
         logger.info("Generation completed successfully.")
@@ -107,14 +112,20 @@ class DiffusionEngine:
             outputs = [outputs] if outputs is not None else []
 
         metrics = {
+            "preprocess_time_ms": preprocess_time * 1000,
+            "diffusion_engine_exec_time_ms": (time.time() - diffusion_engine_start_time) * 1000,
+            "diffusion_engine_total_time_ms": exec_total_time * 1000,
             "image_num": int(request.sampling_params.num_outputs_per_prompt),
             "resolution": int(request.sampling_params.resolution),
-            "postprocess_time_ms": postprocess_time * 1000,
         }
+
         if self.pre_process_func is not None:
             metrics["preprocessing_time_ms"] = preprocess_time * 1000
 
         # Handle single request or multiple requests
+        metrics["postprocess_time_ms"] = postprocess_time * 1000
+        metrics["num_inference_steps"] = int(request.sampling_params.num_inference_steps)
+
         if len(request.prompts) == 1:
             # Single request: return single OmniRequestOutput
             prompt = request.prompts[0]
@@ -201,7 +212,7 @@ class DiffusionEngine:
                         ),
                     )
 
-            return results
+        return results
 
     @staticmethod
     def make_engine(config: OmniDiffusionConfig) -> "DiffusionEngine":