diff --git a/shark/examples/shark_inference/stable_diffusion/stable_args.py b/shark/examples/shark_inference/stable_diffusion/stable_args.py index 03279c96e9..028f13b207 100644 --- a/shark/examples/shark_inference/stable_diffusion/stable_args.py +++ b/shark/examples/shark_inference/stable_diffusion/stable_args.py @@ -4,6 +4,10 @@ description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter ) +############################################################################## +### Stable Diffusion Params +############################################################################## + p.add_argument( "--prompts", nargs="+", @@ -18,10 +22,6 @@ help="text you don't want to see in the generated image.", ) -p.add_argument( - "--device", type=str, default="cpu", help="device to run the model." -) - p.add_argument( "--steps", type=int, @@ -29,13 +29,6 @@ help="the no. of steps to do the sampling.", ) -p.add_argument( - "--version", - type=str, - default="v2.1base", - help="Specify version of stable diffusion model", -) - p.add_argument( "--seed", type=int, @@ -51,10 +44,25 @@ ) p.add_argument( - "--import_mlir", - default=False, - action=argparse.BooleanOptionalAction, - help="imports the model from torch module to shark_module otherwise downloads the model from shark_tank.", + "--max_length", + type=int, + default=77, + help="max length of the tokenizer output.", +) + +############################################################################## +### Model Config and Usage Params +############################################################################## + +p.add_argument( + "--device", type=str, default="cpu", help="device to run the model." +) + +p.add_argument( + "--version", + type=str, + default="v2.1base", + help="Specify version of stable diffusion model", ) p.add_argument( @@ -62,10 +70,10 @@ ) p.add_argument( - "--max_length", - type=int, - default=77, - help="max length of the tokenizer output.", + "--import_mlir", + default=False, + action=argparse.BooleanOptionalAction, + help="imports the model from torch module to shark_module otherwise downloads the model from shark_tank.", ) p.add_argument( @@ -82,6 +90,17 @@ help="saves the compiled flatbuffer to the local directory", ) +p.add_argument( + "--use_tuned", + default=False, + action=argparse.BooleanOptionalAction, + help="Download and use the tuned version of the model if available", +) + +############################################################################## +### IREE - Vulkan supported flags +############################################################################## + p.add_argument( "--iree-vulkan-target-triple", type=str, @@ -97,12 +116,22 @@ ) p.add_argument( - "--use_tuned", + "--vulkan_large_heap_block_size", + default="4294967296", + help="flag for setting VMA preferredLargeHeapBlockSize for vulkan device, default is 4G", +) + +p.add_argument( + "--vulkan_validation_layers", default=False, action=argparse.BooleanOptionalAction, - help="Download and use the tuned version of the model if available", + help="flag for disabling vulkan validation layers when benchmarking", ) +############################################################################## +### Misc. Debug and Optimization flags +############################################################################## + p.add_argument( "--local_tank_cache", default="", @@ -128,17 +157,10 @@ help='directory where you want to store dispatch data generated with "--dispatch_benchmarks"', ) -p.add_argument( - "--vulkan_large_heap_block_size", - default="4294967296", - help="flag for setting VMA preferredLargeHeapBlockSize for vulkan device, default is 4G", -) - p.add_argument( "--enable_rgp", default=False, action=argparse.BooleanOptionalAction, help="flag for inserting debug frames between iterations for use with rgp.", ) - args = p.parse_args() diff --git a/shark/examples/shark_inference/stable_diffusion/utils.py b/shark/examples/shark_inference/stable_diffusion/utils.py index 5dc9d30038..3a4044e217 100644 --- a/shark/examples/shark_inference/stable_diffusion/utils.py +++ b/shark/examples/shark_inference/stable_diffusion/utils.py @@ -75,6 +75,7 @@ def set_iree_runtime_flags(): vulkan_runtime_flags = [ f"--vulkan_large_heap_block_size={args.vulkan_large_heap_block_size}", + f"--vulkan_validation_layers={'true' if args.vulkan_validation_layers else 'false'}", ] if args.enable_rgp: vulkan_runtime_flags += [