Compatible with vllm-omni 0.16.0

Signed-off-by: Chen Yang <2082464740@qq.com>

Author: erfgss

vllm_omni/diffusion/diffusion_engine.py

@@ -67,14 +67,22 @@ def __init__(self, od_config: OmniDiffusionConfig):
             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.")
@@ -102,15 +110,28 @@ def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
         if not isinstance(outputs, list):
             outputs = [outputs] if outputs is not None else []
 
+
         metrics = {
+            "preprocess_time_ms": round(preprocess_time * 1000,2),
+            "diffusion_engine_exec_time_ms": round((time.time() - diffusion_engine_start_time) * 1000,2),
+            "dit_time_ms": round(exec_total_time * 1000,2),
+            "postprocess_time_ms": round(postprocess_time * 1000,2),
             "image_num": int(request.sampling_params.num_outputs_per_prompt),
             "resolution": int(request.sampling_params.resolution),
-            "postprocess_time_ms": postprocess_time * 1000,
+            "denoise_time_per_step_ms": 0.0,
+            "vae_time_ms": 0.0,
         }
-        if self.pre_process_func is not None:
-            metrics["preprocessing_time_ms"] = preprocess_time * 1000
 
-        # Handle single request or multiple requests
+
+        dit_time_seconds = metrics["dit_time_ms"] / 1000
+        num_steps = request.sampling_params.num_inference_steps
+
+        if num_steps > 0:
+            total_denoise_time = dit_time_seconds 
+            metrics["denoise_time_per_step_ms"] = round((total_denoise_time / num_steps) * 1000,2)
+
+        metrics["vae_time_ms"] = round(dit_time_seconds  * 1000,2)
+
         if len(request.prompts) == 1:
             # Single request: return single OmniRequestOutput
             prompt = request.prompts[0]
@@ -177,7 +198,7 @@ def step(self, request: OmniDiffusionRequest) -> list[OmniRequestOutput]:
                         )
                     )
 
-            return results
+        return results
 
     @staticmethod
     def make_engine(config: OmniDiffusionConfig) -> "DiffusionEngine":
@@ -378,4 +399,4 @@ def close(self) -> None:
     def abort(self, request_id: str | Iterable[str]) -> None:
         # TODO implement it
         logger.warning("DiffusionEngine abort is not implemented yet")
-        pass
+        pass

vllm_omni/outputs.py

@@ -235,19 +235,63 @@ def __repr__(self) -> str:
         """Custom repr to properly show image count instead of image objects."""
         # For images, show count instead of full list
         images_repr = f"[{len(self.images)} PIL Images]" if self.images else "[]"
-
         # Build repr string
+
+        def _repr_nested(obj) -> str:
+            if isinstance(obj, list):
+                return "[" + ", ".join(_repr_nested(x) for x in obj) + "]"
+            if isinstance(obj, OmniRequestOutput):
+                return obj._repr_multiline(indent="  ")
+            return repr(obj)
+
         parts = [
             f"request_id={self.request_id!r}",
             f"finished={self.finished}",
             f"stage_id={self.stage_id}",
             f"final_output_type={self.final_output_type!r}",
-            f"request_output={self.request_output}",
+            f"request_output={_repr_nested(self.request_output)}",
             f"images={images_repr}",
             f"prompt={self.prompt!r}",
             f"latents={self.latents}",
             f"metrics={self.metrics}",
             f"multimodal_output={self._multimodal_output}",
         ]
-
         return f"OmniRequestOutput({', '.join(parts)})"
+
+ 
+    def _repr_multiline(self, indent: str = "") -> str:
+        """Helper to produce multi-line, indented repr for nested logging."""
+        images_repr = f"[{len(self.images)} PIL Images]" if self.images else "[]"
+
+        def _repr_nested(obj, ind: str) -> str:
+            if isinstance(obj, list):
+                inner = ",\n".join(_repr_nested(x, ind + "  ") for x in obj)
+                return "[\n" + inner + "\n" + ind + "]"
+            if isinstance(obj, OmniRequestOutput):
+                return obj._repr_multiline(indent=ind + "  ")
+            return repr(obj)
+
+        # Format metrics with each key-value pair on a separate line
+        if self.metrics:
+            metrics_indent = indent + "    "
+            metrics_lines = f",\n{metrics_indent}".join(
+                f"{k!r}: {v!r}" for k, v in self.metrics.items()
+            )
+            metrics_repr = f"{{\n{metrics_indent}{metrics_lines}\n{indent}  }}"
+        else:
+            metrics_repr = "{}"
+
+        lines = [
+            f"{indent}OmniRequestOutput(",
+            f"{indent}  request_id={self.request_id!r},",
+            f"{indent}  finished={self.finished},",
+            f"{indent}  stage_id={self.stage_id},",
+            f"{indent}  final_output_type={self.final_output_type!r},",
+            f"{indent}  request_output={_repr_nested(self.request_output, indent + '  ')},",
+            f"{indent}  images={images_repr},",
+            f"{indent}  prompt={self.prompt!r},",
+            f"{indent}  latents={self.latents},",
+            f"{indent}  metrics={metrics_repr},",
+            f"{indent})",
+        ]
+        return "\n".join(lines)