Add VAE cuda tuned model (huggingface#796)

Author: yzhang93

shark/examples/shark_inference/stable_diffusion/opt_params.py

@@ -79,12 +79,15 @@ def get_unet():
 
 def get_vae():
     # Tuned model is present only for `fp16` precision.
-    is_tuned = (
-        "tuned" if (args.use_tuned and "vulkan" in args.device) else "untuned"
-    )
+    is_tuned = "tuned" if args.use_tuned else "untuned"
     is_base = "/base" if args.use_base_vae else ""
-    bucket_key = f"{args.variant}/{is_tuned}"
-    model_key = f"{args.variant}/{args.version}/vae/{args.precision}/length_77/{is_tuned}{is_base}"
+    if "vulkan" not in args.device and args.use_tuned:
+        bucket_key = f"{args.variant}/{is_tuned}/{args.device}"
+        model_key = f"{args.variant}/{args.version}/vae/{args.precision}/length_77/{is_tuned}{is_base}/{args.device}"
+    else:
+        bucket_key = f"{args.variant}/{is_tuned}"
+        model_key = f"{args.variant}/{args.version}/vae/{args.precision}/length_77/{is_tuned}{is_base}"
+
     bucket, model_name, iree_flags = get_params(
         bucket_key, model_key, "vae", is_tuned, args.precision
     )

shark/examples/shark_inference/stable_diffusion/resources/model_db.json

@@ -27,8 +27,10 @@
     "stablediffusion/v2_1base/unet/fp16/length_64/tuned/cuda":"unet_19dec_v2p1base_fp16_64_cuda_tuned",
     "stablediffusion/v2_1base/vae/fp16/length_77/untuned":"vae2base_19dec_fp16",
     "stablediffusion/v2_1base/vae/fp16/length_77/tuned":"vae2base_19dec_fp16_tuned",
+    "stablediffusion/v2_1base/vae/fp16/length_77/tuned/cuda":"vae2base_19dec_fp16_cuda_tuned",
     "stablediffusion/v2_1base/vae/fp16/length_77/untuned/base":"vae2base_8dec_fp16",
     "stablediffusion/v2_1base/vae/fp16/length_77/tuned/base":"vae2base_8dec_fp16_tuned",
+    "stablediffusion/v2_1base/vae/fp16/length_77/tuned/base/cuda":"vae2base_8dec_fp16_cuda_tuned",
     "stablediffusion/v2_1base/clip/fp32/length_77/untuned":"clip2base_18dec_fp32",
     "stablediffusion/v2_1base/clip/fp32/length_64/untuned":"clip_19dec_v2p1base_fp32_64",
     "stablediffusion/v2_1/unet/fp16/length_77/untuned":"unet2_14dec_fp16",

shark/examples/shark_inference/stable_diffusion/sd_annotation.py

@@ -43,8 +43,8 @@
     winograd_config_dir = f"{WORKDIR}configs/" + config_name
     download_public_file(full_gs_url, winograd_config_dir, True)
 
-if args.annotation_model == "unet":
-    if args.variant in ["anythingv3", "analogdiffusion"]:
+if args.annotation_model == "unet" or device == "cuda":
+    if args.variant in ["anythingv3", "analogdiffusion"] or args.annotation_model == "vae":
         args.max_length = 77
     config_name = f"{args.annotation_model}_{args.version}_{args.precision}_len{args.max_length}_{device}.json"
     full_gs_url = config_bucket + config_name
@@ -67,7 +67,7 @@
             f.write(str(winograd_model))
 
 # For Unet annotate the model with tuned lowering configs
-if args.annotation_model == "unet":
+if args.annotation_model == "unet" or device == "cuda":
     if args.use_winograd:
         input_mlir = f"{args.annotation_output}/{model_name}_tuned_torch.mlir"
         dump_after = "iree-linalg-ext-convert-conv2d-to-winograd"