[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>

examples/offline_inference/custom_pipeline/image_to_image/image_edit.py

@@ -99,6 +99,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--vae_use_slicing", action="store_true")
     parser.add_argument("--vae_use_tiling", action="store_true")
     parser.add_argument("--enable-cpu-offload", action="store_true")
+    parser.add_argument("--log-stats", "--log_stats", dest="log_stats", action="store_true", default=False)
 
     return parser.parse_args()
 
@@ -155,6 +156,7 @@ async def main():
         cache_config=cache_config,
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
+        log_stats=args.log_stats,
         enable_cpu_offload=args.enable_cpu_offload,
         diffusion_load_format="dummy",
         custom_pipeline_args={"pipeline_class": "custom_pipeline.CustomPipeline"},

examples/offline_inference/image_to_image/image_edit.py

@@ -320,6 +320,11 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable layerwise (blockwise) offloading on DiT modules.",
     )
+    parser.add_argument(
+        "--log-stats",
+        action="store_true",
+        help="Enable vLLM-Omni statistics logging.",
+    )
     return parser.parse_args()
 
 
@@ -383,6 +388,7 @@ def main():
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
         enable_cpu_offload=args.enable_cpu_offload,
+        log_stats=args.log_stats,
     )
     print("Pipeline loaded")
 

examples/offline_inference/image_to_video/image_to_video.py

@@ -176,6 +176,11 @@ def parse_args() -> argparse.Namespace:
             "Default 1 means pure sharding (no replication). "
         ),
     )
+    parser.add_argument(
+        "--log-stats",
+        action="store_true",
+        help="Enable vLLM-Omni statistics logging.",
+    )
     return parser.parse_args()
 
 
@@ -278,6 +283,7 @@ def main():
         enable_cpu_offload=args.enable_cpu_offload,
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
+        log_stats=args.log_stats,
         model_class_name=model_class_name,
         cache_backend=args.cache_backend,
         cache_config=cache_config,

examples/offline_inference/text_to_image/text_to_image.py

@@ -94,6 +94,11 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable cache-dit summary logging after diffusion forward passes.",
     )
+    parser.add_argument(
+        "--log-stats",
+        action="store_true",
+        help="Enable vLLM-Omni statistics logging.",
+    )
     parser.add_argument(
         "--ulysses-degree",
         type=int,
@@ -299,6 +304,7 @@ def main():
         "parallel_config": parallel_config,
         "enforce_eager": args.enforce_eager,
         "enable_cpu_offload": args.enable_cpu_offload,
+        "log_stats": args.log_stats,
         **lora_args,
         **quant_kwargs,
     }

examples/offline_inference/text_to_video/text_to_video.py

@@ -127,6 +127,11 @@ def parse_args() -> argparse.Namespace:
         default=1,
         help="Number of GPUs used for VAE patch/tile parallelism (decode).",
     )
+    parser.add_argument(
+        "--log-stats",
+        action="store_true",
+        help="Enable vLLM-Omni statistics logging.",
+    )
     return parser.parse_args()
 
 
@@ -176,6 +181,7 @@ def main():
         enable_cpu_offload=args.enable_cpu_offload,
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
+        log_stats=args.log_stats,
         cache_backend=args.cache_backend,
         cache_config=cache_config,
     )