Add CUDA tuned models for SD variants (huggingface#814)

Author: yzhang93

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

@@ -5,23 +5,28 @@
     "stablediffusion/tuned/cuda":"gs://shark_tank/sd_tuned/cuda",
     "anythingv3/untuned":"gs://shark_tank/sd_anythingv3",
     "anythingv3/tuned":"gs://shark_tank/sd_tuned",
+    "anythingv3/tuned/cuda":"gs://shark_tank/sd_tuned/cuda",
     "analogdiffusion/untuned":"gs://shark_tank/sd_analog_diffusion",
     "analogdiffusion/tuned":"gs://shark_tank/sd_tuned",
+    "analogdiffusion/tuned/cuda":"gs://shark_tank/sd_tuned/cuda",
     "openjourney/untuned":"gs://shark_tank/sd_openjourney",
     "openjourney/tuned":"gs://shark_tank/sd_tuned",
     "dreamlike/untuned":"gs://shark_tank/sd_dreamlike_diffusion"
   },
   {
     "stablediffusion/v1_4/unet/fp16/length_77/untuned":"unet_8dec_fp16",
     "stablediffusion/v1_4/unet/fp16/length_77/tuned":"unet_8dec_fp16_tuned",
+    "stablediffusion/v1_4/unet/fp16/length_77/tuned/cuda":"unet_8dec_fp16_cuda_tuned",
     "stablediffusion/v1_4/unet/fp32/length_77/untuned":"unet_1dec_fp32",
     "stablediffusion/v1_4/vae/fp16/length_77/untuned":"vae_19dec_fp16",
     "stablediffusion/v1_4/vae/fp16/length_77/tuned":"vae_19dec_fp16_tuned",
+    "stablediffusion/v1_4/vae/fp16/length_77/tuned/cuda":"vae_19dec_fp16_cuda_tuned",
     "stablediffusion/v1_4/vae/fp16/length_77/untuned/base":"vae_8dec_fp16",
     "stablediffusion/v1_4/vae/fp32/length_77/untuned":"vae_1dec_fp32",
     "stablediffusion/v1_4/clip/fp32/length_77/untuned":"clip_18dec_fp32",
     "stablediffusion/v2_1base/unet/fp16/length_77/untuned":"unet2base_8dec_fp16",
     "stablediffusion/v2_1base/unet/fp16/length_77/tuned":"unet2base_8dec_fp16_tuned_v2",
+    "stablediffusion/v2_1base/unet/fp16/length_77/tuned/cuda":"unet2base_8dec_fp16_cuda_tuned",
     "stablediffusion/v2_1base/unet/fp16/length_64/untuned":"unet_19dec_v2p1base_fp16_64",
     "stablediffusion/v2_1base/unet/fp16/length_64/tuned":"unet_19dec_v2p1base_fp16_64_tuned",
     "stablediffusion/v2_1base/unet/fp16/length_64/tuned/cuda":"unet_19dec_v2p1base_fp16_64_cuda_tuned",
@@ -39,18 +44,22 @@
     "stablediffusion/v2_1/clip/fp32/length_77/untuned":"clip2_18dec_fp32",
     "anythingv3/v2_1base/unet/fp16/length_77/untuned":"av3_unet_19dec_fp16",
     "anythingv3/v2_1base/unet/fp16/length_77/tuned":"av3_unet_19dec_fp16_tuned",
+    "anythingv3/v2_1base/unet/fp16/length_77/tuned/cuda":"av3_unet_19dec_fp16_cuda_tuned",
     "anythingv3/v2_1base/unet/fp32/length_77/untuned":"av3_unet_19dec_fp32",
     "anythingv3/v2_1base/vae/fp16/length_77/untuned":"av3_vae_19dec_fp16",
     "anythingv3/v2_1base/vae/fp16/length_77/tuned":"av3_vae_19dec_fp16_tuned",
+    "anythingv3/v2_1base/vae/fp16/length_77/tuned/cuda":"av3_vae_19dec_fp16_cuda_tuned",
     "anythingv3/v2_1base/vae/fp16/length_77/untuned/base":"av3_vaebase_22dec_fp16",
     "anythingv3/v2_1base/vae/fp32/length_77/untuned":"av3_vae_19dec_fp32",
     "anythingv3/v2_1base/vae/fp32/length_77/untuned/base":"av3_vaebase_22dec_fp32",
     "anythingv3/v2_1base/clip/fp32/length_77/untuned":"av3_clip_19dec_fp32",
     "analogdiffusion/v2_1base/unet/fp16/length_77/untuned":"ad_unet_19dec_fp16",
     "analogdiffusion/v2_1base/unet/fp16/length_77/tuned":"ad_unet_19dec_fp16_tuned",
+    "analogdiffusion/v2_1base/unet/fp16/length_77/tuned/cuda":"ad_unet_19dec_fp16_cuda_tuned",
     "analogdiffusion/v2_1base/unet/fp32/length_77/untuned":"ad_unet_19dec_fp32",
     "analogdiffusion/v2_1base/vae/fp16/length_77/untuned":"ad_vae_19dec_fp16",
     "analogdiffusion/v2_1base/vae/fp16/length_77/tuned":"ad_vae_19dec_fp16_tuned",
+    "analogdiffusion/v2_1base/vae/fp16/length_77/tuned/cuda":"ad_vae_19dec_fp16_cuda_tuned",
     "analogdiffusion/v2_1base/vae/fp16/length_77/untuned/base":"ad_vaebase_22dec_fp16",
     "analogdiffusion/v2_1base/vae/fp32/length_77/untuned":"ad_vae_19dec_fp32",
     "analogdiffusion/v2_1base/vae/fp32/length_77/untuned/base":"ad_vaebase_22dec_fp32",

shark/examples/shark_inference/stable_diffusion/sd_annotation.py

@@ -44,10 +44,10 @@
     download_public_file(full_gs_url, winograd_config_dir, True)
 
 if args.annotation_model == "unet" or device == "cuda":
-    if (
-        args.variant in ["anythingv3", "analogdiffusion"]
-        or args.annotation_model == "vae"
-    ):
+    if args.variant in ["anythingv3", "analogdiffusion"]:
+        args.max_length = 77
+        args.version = "v1_4"
+    if 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

shark/examples/shark_inference/stable_diffusion/utils.py

@@ -205,11 +205,10 @@ def set_init_device_flags():
 
     # Use tuned model in the case of stablediffusion/fp16 and cuda device sm_80
     if (
-        args.variant == "stablediffusion"
+        args.variant in ["stablediffusion", "anythingv3", "analogdiffusion"]
         and args.precision == "fp16"
         and "cuda" in args.device
         and get_cuda_sm_cc() == "sm_80"
-        and args.version == "v2_1base"
     ):
         args.use_tuned = True