Tweak compile-time flags for SD submodels.

nod-ai · Feb 20, 2024 · 60c013e · 60c013e
1 parent 92c11be
commit 60c013e
Showing 1 changed file with 9 additions and 11 deletions.
diff --git a/apps/shark_studio/api/sd.py b/apps/shark_studio/api/sd.py
@@ -45,16 +45,15 @@
             "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-pad-linalg-ops{pad-size=16}))",
         ],
     },
-    # "vae_encode": {
-    #     "initializer": vae.export_vae_model,
-    #     "ireec_flags": [
-    #         "--iree-flow-collapse-reduction-dims",
-    #         "--iree-opt-const-expr-hoisting=False",
-    #         "--iree-codegen-linalg-max-constant-fold-elements=9223372036854775807",
-    #         "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-global-opt-detach-elementwise-from-named-ops,iree-global-opt-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32},iree-linalg-ext-convert-conv2d-to-winograd))",
-    #         "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-pad-linalg-ops{pad-size=16}))",
-    #     ],
-    # },
+    "vae_encode": {
+        "initializer": vae.export_vae_model,
+        "ireec_flags": [
+            "--iree-flow-collapse-reduction-dims",
+            "--iree-opt-const-expr-hoisting=False",
+            "--iree-codegen-linalg-max-constant-fold-elements=9223372036854775807",
+            "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-global-opt-detach-elementwise-from-named-ops,iree-global-opt-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32},iree-linalg-ext-convert-conv2d-to-winograd))",
+        ],
+    },
     "unet": {
         "initializer": unet.export_unet_model,
         "ireec_flags": [
@@ -71,7 +70,6 @@
             "--iree-opt-const-expr-hoisting=False",
             "--iree-codegen-linalg-max-constant-fold-elements=9223372036854775807",
             "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-global-opt-detach-elementwise-from-named-ops,iree-global-opt-convert-1x1-filter-conv2d-to-matmul,iree-preprocessing-convert-conv2d-to-img2col,iree-preprocessing-pad-linalg-ops{pad-size=32},iree-linalg-ext-convert-conv2d-to-winograd))",
-            "--iree-preprocessing-pass-pipeline=builtin.module(func.func(iree-preprocessing-pad-linalg-ops{pad-size=16}))",
         ],
     },
 }