[skip ci] Add width and height args to offline i2i example script (#4031)

Signed-off-by: Huang, Zeyu <11222265+fhfuih@users.noreply.github.com>

examples/offline_inference/image_to_image/image_edit.py

@@ -143,6 +143,20 @@ def parse_args() -> argparse.Namespace:
         default=None,
         required=False,
     )
+    parser.add_argument(
+        "--width",
+        type=int,
+        default=None,
+        metavar="W",
+        help="Output image width in pixels. Default: None (pipeline's default).",
+    )
+    parser.add_argument(
+        "--height",
+        type=int,
+        default=None,
+        metavar="H",
+        help="Output image height in pixels. Default: None (pipeline's default).",
+    )
     parser.add_argument(
         "--seed",
         type=int,
@@ -236,8 +250,8 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--resolution",
         type=int,
-        default=640,
-        help="Bucket in (640, 1024) to determine the condition and output resolution",
+        default=None,
+        help="Bucket in (640, 1024) to determine the condition and output resolution. If width and height are not provided, this will be set to default 640.",
     )
 
     parser.add_argument(
@@ -384,6 +398,15 @@ def parse_args() -> argparse.Namespace:
 def main():
     args = parse_args()
 
+    if args.resolution and (args.width or args.height):
+        raise ValueError("--resolution and --width/--height cannot be specified together")
+    if args.width is not None and args.width <= 0:
+        raise ValueError("--width must be a positive integer")
+    if args.height is not None and args.height <= 0:
+        raise ValueError("--height must be a positive integer")
+    if not args.width and not args.height and not args.resolution:
+        args.resolution = 640
+
     # Validate input images exist and load them
     input_images = []
     for image_path in args.image:
@@ -464,6 +487,7 @@ def main():
     print(
         f"  Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size}, enable_expert_parallel: {args.enable_expert_parallel}"
     )
+    print(f"  Output image size: {args.width or 'default'}x{args.height or 'default'}")
     print(f"{'=' * 60}\n")
 
     generation_start = time.perf_counter()
@@ -473,22 +497,26 @@ def main():
         omni.start_profile()
 
     # Generate edited image
+    sampling_params = OmniDiffusionSamplingParams(
+        width=args.width,
+        height=args.height,
+        generator=generator,
+        true_cfg_scale=args.cfg_scale,
+        guidance_scale=args.guidance_scale,
+        guidance_scale_2=args.guidance_scale_2,
+        num_inference_steps=args.num_inference_steps,
+        num_outputs_per_prompt=args.num_outputs_per_prompt,
+        layers=args.layers,
+    )
+    if args.resolution:
+        sampling_params.resolution = args.resolution
     outputs = omni.generate(
         {
             "prompt": args.prompt,
             "negative_prompt": args.negative_prompt,
             "multi_modal_data": {"image": input_image},
         },
-        OmniDiffusionSamplingParams(
-            generator=generator,
-            true_cfg_scale=args.cfg_scale,
-            guidance_scale=args.guidance_scale,
-            guidance_scale_2=args.guidance_scale_2,
-            num_inference_steps=args.num_inference_steps,
-            num_outputs_per_prompt=args.num_outputs_per_prompt,
-            layers=args.layers,
-            resolution=args.resolution,
-        ),
+        sampling_params,
     )
     generation_end = time.perf_counter()
     generation_time = generation_end - generation_start