Add the stable diffusion v2.1 version.

Author: Prashant Kumar

shark/examples/shark_inference/stable_diffusion/main.py

@@ -15,7 +15,6 @@
 from opt_params import get_unet, get_vae, get_clip
 import time
 import sys
-from model_wrappers import get_vae_mlir
 from shark.iree_utils.compile_utils import dump_isas
 
 # Helper function to profile the vulkan device.
@@ -43,7 +42,7 @@ def end_profiling(device):
     neg_prompt = args.negative_prompts
     height = 512  # default height of Stable Diffusion
     width = 512  # default width of Stable Diffusion
-    if args.version == "v2":
+    if args.version == "v2.1":
         height = 768
         width = 768
 
@@ -75,13 +74,13 @@ def end_profiling(device):
         "CompVis/stable-diffusion-v1-4",
         subfolder="scheduler",
     )
-    if args.version == "v2":
+    if args.version == "v2.1":
         tokenizer = CLIPTokenizer.from_pretrained(
-            "stabilityai/stable-diffusion-2", subfolder="tokenizer"
+            "stabilityai/stable-diffusion-2-1", subfolder="tokenizer"
         )
 
         scheduler = DPMSolverMultistepScheduler.from_pretrained(
-            "stabilityai/stable-diffusion-2",
+            "stabilityai/stable-diffusion-2-1",
             subfolder="scheduler",
         )
 

shark/examples/shark_inference/stable_diffusion/model_wrappers.py

@@ -5,13 +5,13 @@
 import torch
 
 model_config = {
-    "v2": "stabilityai/stable-diffusion-2",
+    "v2.1": "stabilityai/stable-diffusion-2-1",
     "v2.1base": "stabilityai/stable-diffusion-2-1-base",
     "v1.4": "CompVis/stable-diffusion-v1-4",
 }
 
 model_input = {
-    "v2": {
+    "v2.1": {
         "clip": (torch.randint(1, 2, (1, 77)),),
         "vae": (torch.randn(1, 4, 96, 96),),
         "unet": (
@@ -52,7 +52,7 @@ def get_clip_mlir(model_name="clip_text", extra_args=[]):
     text_encoder = CLIPTextModel.from_pretrained(
         "openai/clip-vit-large-patch14"
     )
-    if args.version == "v2":
+    if args.version != "v1.4":
         text_encoder = CLIPTextModel.from_pretrained(
             model_config[args.version], subfolder="text_encoder"
         )

shark/examples/shark_inference/stable_diffusion/opt_params.py

@@ -30,6 +30,8 @@ def get_unet():
             model_name = "unet_8dec_fp16"
             if args.version == "v2.1base":
                 model_name = "unet2base_8dec_fp16"
+            if args.version == "v2.1":
+                model_name = "unet2_14dec_fp16"
             iree_flags += [
                 "--iree-flow-enable-padding-linalg-ops",
                 "--iree-flow-linalg-ops-padding-size=32",
@@ -79,6 +81,8 @@ def get_vae():
         model_name = "vae_8dec_fp16"
         if args.version == "v2.1base":
             model_name = "vae2base_8dec_fp16"
+        if args.version == "v2.1":
+            model_name = "vae2_14dec_fp16"
         iree_flags += [
             "--iree-flow-enable-padding-linalg-ops",
             "--iree-flow-linalg-ops-padding-size=32",
@@ -144,6 +148,8 @@ def get_clip():
     model_name = "clip_8dec_fp32"
     if args.version == "v2.1base":
         model_name = "clip2base_8dec_fp32"
+    if args.version == "v2.1":
+        model_name = "clip2_14dec_fp32"
     iree_flags += [
         "--iree-flow-linalg-ops-padding-size=16",
         "--iree-flow-enable-padding-linalg-ops",