[Diffusion] Benchmark tweaks (#48)

Author: jaywonchung

configs/xdit/text-to-image/benchmark.yaml

@@ -9,7 +9,8 @@ command_template: |
     --workload.base_dir run/diffusion/text-to-image/{model_id} \
     --workload.batch_size {batch_size} \
     --workload.ulysses_degree {ulysses_degree} \
-    --workload.ring_degree {ring_degree}
+    --workload.ring_degree {ring_degree} \
+    --workload.inference_steps None
 
 sweep_defaults:
   - batch_size: [1, 2, 4, 8, 16]

configs/xdit/text-to-video/benchmark.yaml

@@ -10,6 +10,8 @@ command_template: |
     --workload.batch_size {batch_size} \
     --workload.ulysses_degree {ulysses_degree} \
     --workload.ring_degree {ring_degree}
+    --workload.inference_steps None \
+    --workload.num_frames None
 
 sweep_defaults:
   - batch_size: [1, 2, 4, 8, 16]
@@ -42,5 +44,3 @@ sweep_defaults:
   - batch_size: [1, 2, 4, 8, 16]
     ulysses_degree: [8]
     ring_degree: [1]
-
-

mlenergy/diffusion/benchmark.py

@@ -965,6 +965,12 @@ def save_results(
     result_json["inference_steps"] = args.workload.inference_steps
     result_json["seed"] = args.workload.seed
 
+    # Video-specific parameters
+    if hasattr(args.workload, "num_frames") and args.workload.num_frames is not None:
+        result_json["num_frames"] = args.workload.num_frames
+    if hasattr(args.workload, "fps") and args.workload.fps is not None:
+        result_json["fps"] = args.workload.fps
+
     # Performance metrics
     result_json["total_images"] = num_images
     result_json["total_time"] = benchmark_duration