[SD][WEB] Add vae tuned model in the SD web (#653)

1. Add tuned vae model in the SD web. 2. Use tuned models in case of rdna3 cards. Signed-Off-by: Gaurav Shukla <gaurav@nod-labs.com> Signed-off-by: Gaurav Shukla <gaurav@nod-labs.com>
nod-ai · Dec 16, 2022 · 10160a0 · 10160a0
1 parent 72976a2
commit 10160a0
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 14 deletions.
diff --git a/web/models/stable_diffusion/cache_objects.py b/web/models/stable_diffusion/cache_objects.py
@@ -40,6 +40,10 @@
     subfolder="scheduler",
 )
 
+# use tuned unet model in case of rdna3 cards.
+if "rdna3" in get_vulkan_triple_flag():
+    args.use_tuned = True
+
 # set iree-runtime flags
 set_iree_runtime_flags()
 

diff --git a/web/models/stable_diffusion/opt_params.py b/web/models/stable_diffusion/opt_params.py
@@ -86,20 +86,32 @@ def get_vae():
     if sys.platform == "darwin":
         iree_flags.append("-iree-stream-fuse-binding=false")
     if args.precision in ["fp16", "int8"]:
-        bucket = "gs://shark_tank/stable_diffusion"
-        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",
-            "--iree-flow-enable-conv-img2col-transform",
-        ]
-        if args.import_mlir:
-            return get_vae_mlir(model_name, iree_flags)
-        return get_shark_model(bucket, model_name, iree_flags)
+        if args.use_tuned:
+            bucket = "gs://shark_tank/vivian"
+            if args.version == "v2.1base":
+                model_name = "vae2base_8dec_fp16_tuned"
+            iree_flags += [
+                "--iree-flow-enable-padding-linalg-ops",
+                "--iree-flow-linalg-ops-padding-size=32",
+                "--iree-flow-enable-conv-img2col-transform",
+                "--iree-flow-enable-conv-winograd-transform",
+            ]
+            return get_shark_model(bucket, model_name, iree_flags)
+        else:
+            bucket = "gs://shark_tank/stable_diffusion"
+            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",
+                "--iree-flow-enable-conv-img2col-transform",
+            ]
+            if args.import_mlir:
+                return get_vae_mlir(model_name, iree_flags)
+            return get_shark_model(bucket, model_name, iree_flags)
 
     if args.precision == "fp32":
         bucket = "gs://shark_tank/stable_diffusion"