Decompose VAE for cpu

nod-ai · Oct 4, 2024 · fc6d018 · fc6d018
1 parent e630d39
commit fc6d018
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/models/turbine_models/tests/sdxl_test.py b/models/turbine_models/tests/sdxl_test.py
@@ -93,7 +93,11 @@ def test00_sdxl_pipe(self):
         decomp_attn = {
             "text_encoder": True,
             "unet": False,
-            "vae": False,
+            "vae": (
+                False
+                if any(x in arguments["device"] for x in ["hip", "rocm"])
+                else True
+            ),
         }
         self.pipe = SharkSDPipeline(
             arguments["hf_model_name"],
@@ -377,7 +381,7 @@ def test04_ExportVaeModelDecode(self):
                     "bs" + str(arguments["batch_size"]),
                     str(arguments["height"]) + "x" + str(arguments["width"]),
                     arguments["precision"],
-                    "vae",
+                    "vae" if arguments["device"] != "cpu" else "vae_decomp_attn",
                     arguments["iree_target_triple"],
                 ]
             )