From 780f520f02e25f5ce144e82cbdc0da87509371e6 Mon Sep 17 00:00:00 2001
From: Ean Garvey <87458719+monorimet@users.noreply.github.com>
Date: Mon, 11 Sep 2023 13:42:52 -0500
Subject: [PATCH] Fix vk.target_env extensions and remove redundant SD imports.
 (#1826)

* Remove redundant IREE runtime imports.

* Fix vulkan target env extensions.
---
 apps/stable_diffusion/shark_studio_imports.py | 1 -
 shark/iree_utils/vulkan_target_env_utils.py   | 3 ++-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/apps/stable_diffusion/shark_studio_imports.py b/apps/stable_diffusion/shark_studio_imports.py
index e4ba324976..4a55a94bf1 100644
--- a/apps/stable_diffusion/shark_studio_imports.py
+++ b/apps/stable_diffusion/shark_studio_imports.py
@@ -81,4 +81,3 @@
     if not any(kw in x for kw in blacklist)
 ]
 hiddenimports += [x for x in collect_submodules("iree") if "tests" not in x]
-hiddenimports += ["iree._runtime", "iree._runtime_libs"]
diff --git a/shark/iree_utils/vulkan_target_env_utils.py b/shark/iree_utils/vulkan_target_env_utils.py
index d92dd6d9f2..e1de604521 100644
--- a/shark/iree_utils/vulkan_target_env_utils.py
+++ b/shark/iree_utils/vulkan_target_env_utils.py
@@ -57,7 +57,8 @@ def get_version(triple):
 @functools.cache
 def get_extensions(triple):
     def make_ext_list(ext_list):
-        return ", ".join(ext_list)
+        res = ", ".join(ext_list)
+        return f"[{res}]"
 
     arch, product, os = triple
     if arch == "m1":