Make separate function for each model.

Author: Prashant Kumar

shark/examples/shark_inference/stable_diffusion/main.py

@@ -5,41 +5,10 @@
 from tqdm.auto import tqdm
 import numpy as np
 from stable_args import args
-from model_wrappers import (
-    get_vae32,
-    get_vae16,
-    get_unet16_wrapped,
-    get_unet32_wrapped,
-    get_clipped_text,
-)
 from utils import get_shark_model
+from opt_params import get_unet, get_vae, get_clip
 import time
 
-GCLOUD_BUCKET = "gs://shark_tank/prashant_nod"
-VAE_FP16 = "vae_fp16"
-VAE_FP32 = "vae_fp32"
-UNET_FP16 = "unet_fp16"
-UNET_FP32 = "unet_fp32"
-IREE_EXTRA_ARGS = []
-
-TUNED_GCLOUD_BUCKET = "gs://shark_tank/quinn"
-UNET_FP16_TUNED = "unet_fp16_tunedv2"
-
-BATCH_SIZE = len(args.prompts)
-
-if BATCH_SIZE not in [1, 2]:
-    import sys
-
-    sys.exit("Only batch size 1 and 2 are supported.")
-
-if BATCH_SIZE > 1 and args.precision != "fp16":
-    sys.exit("batch size > 1 is supported for fp16 model.")
-
-
-if BATCH_SIZE != 1:
-    TUNED_GCLOUD_BUCKET = "gs://shark_tank/prashant_nod"
-    UNET_FP16_TUNED = f"unet_fp16_{BATCH_SIZE}"
-    VAE_FP16 = f"vae_fp16_{BATCH_SIZE}"
 
 # Helper function to profile the vulkan device.
 def start_profiling(file_path="foo.rdc", profiling_mode="queue"):
@@ -58,87 +27,9 @@ def end_profiling(device):
         return device.end_profiling()
 
 
-def get_models():
-    global IREE_EXTRA_ARGS
-    if args.precision == "fp16":
-        IREE_EXTRA_ARGS += [
-            "--iree-flow-enable-padding-linalg-ops",
-            "--iree-flow-linalg-ops-padding-size=32",
-        ]
-        if args.use_tuned:
-            unet_gcloud_bucket = TUNED_GCLOUD_BUCKET
-            vae_gcloud_bucket = GCLOUD_BUCKET
-            unet_args = IREE_EXTRA_ARGS
-            vae_args = IREE_EXTRA_ARGS + [
-                "--iree-flow-enable-conv-nchw-to-nhwc-transform"
-            ]
-            unet_name = UNET_FP16_TUNED
-            vae_name = VAE_FP16
-        else:
-            unet_gcloud_bucket = GCLOUD_BUCKET
-            vae_gcloud_bucket = GCLOUD_BUCKET
-            IREE_EXTRA_ARGS += [
-                "--iree-flow-enable-conv-nchw-to-nhwc-transform"
-            ]
-            unet_args = IREE_EXTRA_ARGS
-            vae_args = IREE_EXTRA_ARGS
-            unet_name = UNET_FP16
-            vae_name = VAE_FP16
-
-        if batch_size > 1:
-            vae_args = []
-
-        if args.import_mlir == True:
-            return get_vae16(model_name=VAE_FP16), get_unet16_wrapped(
-                model_name=UNET_FP16
-            )
-        else:
-            return get_shark_model(
-                vae_gcloud_bucket,
-                vae_name,
-                vae_args,
-            ), get_shark_model(
-                unet_gcloud_bucket,
-                unet_name,
-                unet_args,
-            )
-
-    elif args.precision == "fp32":
-        IREE_EXTRA_ARGS += [
-            "--iree-flow-enable-conv-nchw-to-nhwc-transform",
-            "--iree-flow-enable-padding-linalg-ops",
-            "--iree-flow-linalg-ops-padding-size=16",
-        ]
-        if args.import_mlir == True:
-            return get_vae32(model_name=VAE_FP32), get_unet32_wrapped(
-                model_name=UNET_FP32
-            )
-        else:
-            return get_shark_model(
-                GCLOUD_BUCKET,
-                VAE_FP32,
-                IREE_EXTRA_ARGS,
-            ), get_shark_model(
-                GCLOUD_BUCKET,
-                UNET_FP32,
-                IREE_EXTRA_ARGS,
-            )
-
-
 if __name__ == "__main__":
 
     dtype = torch.float32 if args.precision == "fp32" else torch.half
-    if len(args.iree_vulkan_target_triple) > 0:
-        IREE_EXTRA_ARGS.append(
-            f"-iree-vulkan-target-triple={args.iree_vulkan_target_triple}"
-        )
-
-    clip_model = "clip_text"
-    clip_extra_args = [
-        "--iree-flow-linalg-ops-padding-size=16",
-        "--iree-flow-enable-padding-linalg-ops",
-    ]
-    clip = get_shark_model(GCLOUD_BUCKET, clip_model, clip_extra_args)
 
     prompt = args.prompts
     height = 512  # default height of Stable Diffusion
@@ -154,7 +45,7 @@ def get_models():
 
     batch_size = len(prompt)
 
-    vae, unet = get_models()
+    unet, vae, clip = get_unet(), get_vae(), get_clip()
 
     tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-large-patch14")
 

shark/examples/shark_inference/stable_diffusion/opt_params.py

@@ -0,0 +1,104 @@
+from model_wrappers import (
+    get_vae32,
+    get_vae16,
+    get_unet16_wrapped,
+    get_unet32_wrapped,
+    get_clipped_text,
+)
+from stable_args import args
+from utils import get_shark_model
+
+BATCH_SIZE = len(args.prompts)
+if BATCH_SIZE != 1:
+    import sys
+
+    sys.exit("Only batch size 1 is supported.")
+
+
+def get_unet():
+    iree_flags = []
+    if len(args.iree_vulkan_target_triple) > 0:
+        iree_flags.append(
+            f"-iree-vulkan-target-triple={args.iree_vulkan_target_triple}"
+        )
+    # Tuned model is present for `fp16` precision.
+    if args.precision == "fp16":
+        if args.use_tuned:
+            bucket = "gs://shark_tank/quinn"
+            model_name = "unet_fp16_tunedv2"
+            iree_flags += [
+                "--iree-flow-enable-padding-linalg-ops",
+                "--iree-flow-linalg-ops-padding-size=32",
+            ]
+            # TODO: Pass iree_flags to the exported model.
+            if args.import_mlir:
+                return get_unet16_wrapped(model_name=model_name)
+            return get_shark_model(bucket, model_name, iree_flags)
+        else:
+            bucket = "gs://shark_tank/prashant_nod"
+            model_name = "unet_fp16"
+            iree_flags += ["--iree-flow-enable-conv-nchw-to-nhwc-transform"]
+            if args.import_mlir:
+                return get_unet16_wrapped(model_name=model_name)
+            return get_shark_model(bucket, model_name, iree_flags)
+
+    # Tuned model is not present for `fp32` case.
+    if args.precision == "fp32":
+        bucket = "gs://shark_tank/prashant_nod"
+        model_name = "unet_fp32"
+        iree_flags += [
+            "--iree-flow-enable-conv-nchw-to-nhwc-transform",
+            "--iree-flow-enable-padding-linalg-ops",
+            "--iree-flow-linalg-ops-padding-size=16",
+        ]
+        if args.import_mlir:
+            return get_unet32_wrapped(model_name=model_name)
+        return get_shark_model(bucket, model_name, iree_flags)
+
+
+def get_vae():
+    iree_flags = []
+    if len(args.iree_vulkan_target_triple) > 0:
+        iree_flags.append(
+            f"-iree-vulkan-target-triple={args.iree_vulkan_target_triple}"
+        )
+    if args.precision == "fp16":
+        bucket = "gs://shark_tank/prashant_nod"
+        model_name = "vae_fp16"
+        iree_flags += [
+            "--iree-flow-enable-conv-nchw-to-nhwc-transform",
+            "--iree-flow-enable-padding-linalg-ops",
+            "--iree-flow-linalg-ops-padding-size=32",
+        ]
+        if args.import_mlir:
+            return get_vae16(model_name)
+        return get_shark_model(bucket, model_name, iree_flags)
+
+    if args.precision == "fp32":
+        bucket = "gs://shark_tank/prashant_nod"
+        model_name = "vae_fp32"
+        iree_flags += [
+            "--iree-flow-enable-conv-nchw-to-nhwc-transform",
+            "--iree-flow-enable-padding-linalg-ops",
+            "--iree-flow-linalg-ops-padding-size=16",
+        ]
+        if args.import_mlir:
+            return get_vae32(model_name)
+        return get_shark_model(bucket, model_name, iree_flags)
+
+
+def get_clip():
+    iree_flags = []
+    if len(args.iree_vulkan_target_triple) > 0:
+        iree_flags.append(
+            f"-iree-vulkan-target-triple={args.iree_vulkan_target_triple}"
+        )
+    bucket = "gs://shark_tank/prashant_nod"
+    model_name = "clip_text"
+    iree_flags = [
+        "--iree-flow-linalg-ops-padding-size=16",
+        "--iree-flow-enable-padding-linalg-ops",
+    ]
+    if args.import_mlir:
+        return get_clipped_text(model_name)
+    return get_shark_model(bucket, model_name, iree_flags)

shark/examples/shark_inference/stable_diffusion/utils.py

@@ -6,7 +6,9 @@
 from torch.fx.experimental.proxy_tensor import make_fx
 from stable_args import args
 from torch._decomp import get_decompositions
-import torch_mlir
+
+if args.import_mlir:
+    import torch_mlir
 
 
 def _compile_module(shark_module, model_name, extra_args=[]):