fix formatting and disable explicit vulkan env settings.

apps/shark_studio/api/initializers.py

@@ -84,7 +84,7 @@ def dumpstacks():
             if line:
                 code.append("  " + line.strip())
     with open(os.path.join(shark_tmp, "stack_dump.log"), "w") as f:
-        f.write("\n".join(code)) 
+        f.write("\n".join(code))
 
 
 def setup_middleware(app):

apps/shark_studio/api/llm.py

@@ -326,6 +326,7 @@ def chat_hf(self, prompt):
         self.global_iter += 1
         return result_output, total_time
 
+
 def llm_chat_api(InputData: dict):
     from datetime import datetime as dt
 
@@ -392,7 +393,6 @@ def llm_chat_api(InputData: dict):
     print("prompt = ", prompt)
 
     for res_op, _ in llm_model.chat(prompt):
-
         if is_chat_completion_api:
             choices = [
                 {
@@ -421,6 +421,7 @@ def llm_chat_api(InputData: dict):
         "choices": choices,
     }
 
+
 if __name__ == "__main__":
     lm = LanguageModel(
         "Trelis/Llama-2-7b-chat-hf-function-calling-v2",

apps/shark_studio/api/sd.py

@@ -83,9 +83,7 @@ def __init__(
             "clip": {"hf_model_name": base_model_id},
             "unet": {
                 "hf_model_name": base_model_id,
-                "unet_model": unet.UnetModel(
-                    hf_model_name=base_model_id
-                ),
+                "unet_model": unet.UnetModel(hf_model_name=base_model_id),
                 "batch_size": batch_size,
                 # "is_controlled": is_controlled,
                 # "num_loras": num_loras,

apps/shark_studio/api/utils.py

@@ -77,6 +77,7 @@ def get_devices_by_name(driver_name):
     available_devices.extend(cpu_device)
     return available_devices
 
+
 def set_init_device_flags():
     if "vulkan" in cmd_opts.device:
         # set runtime flags for vulkan.
@@ -109,6 +110,7 @@ def set_init_device_flags():
     elif "cpu" in cmd_opts.device:
         cmd_opts.device = "cpu"
 
+
 def set_iree_runtime_flags():
     # TODO: This function should be device-agnostic and piped properly
     # to general runtime driver init.
@@ -177,6 +179,7 @@ def get_output_value(dev_dict):
         device_map[f"{driver}://{device['path']}"] = get_output_value(device)
     return device_map
 
+
 def get_opt_flags(model, precision="fp16"):
     iree_flags = []
     if len(cmd_opts.iree_vulkan_target_triple) > 0:
@@ -202,6 +205,7 @@ def get_opt_flags(model, precision="fp16"):
         iree_flags += ["--iree-flow-collapse-reduction-dims"]
     return iree_flags
 
+
 def map_device_to_name_path(device, key_combination=3):
     """Gives the appropriate device data (supported name/path) for user
         selected execution device

apps/shark_studio/modules/pipeline.py

@@ -90,6 +90,8 @@ def get_compiled_map(self, pipe_id, submodel="None", init_kwargs={}) -> None:
                 )
 
                 weights_path = self.get_io_params(submodel)
+                if weights_path:
+                    ireec_flags.append("--iree-opt-const-eval=False")
 
                 self.iree_module_dict[submodel] = get_iree_compiled_module(
                     self.tempfiles[submodel],

apps/shark_studio/studio_imports.py

@@ -56,9 +56,7 @@
 # hidden imports for pyinstaller
 hiddenimports = ["shark", "apps"]
 hiddenimports += [x for x in collect_submodules("gradio") if "tests" not in x]
-hiddenimports += [
-    x for x in collect_submodules("diffusers") if "tests" not in x
-]
+hiddenimports += [x for x in collect_submodules("diffusers") if "tests" not in x]
 blacklist = ["tests", "convert"]
 hiddenimports += [
     x
@@ -67,4 +65,4 @@
 ]
 hiddenimports += [x for x in collect_submodules("iree") if "test" not in x]
 hiddenimports += ["iree._runtime"]
-hiddenimports += [x for x in collect_submodules("scipy") if "test" not in x]
+hiddenimports += [x for x in collect_submodules("scipy") if "test" not in x]

apps/shark_studio/tests/rest_api_test.py

@@ -38,7 +38,6 @@ def llm_chat_test(verbose=False):
 
 
 if __name__ == "__main__":
-
     # "Exercises the chatbot REST API of Shark. Make sure "
     # "Shark is running in API mode on 127.0.0.1:8080 before running"
     # "this script."

apps/shark_studio/web/api/compat.py

@@ -26,7 +26,6 @@
 
 def decode_base64_to_image(encoding):
     if encoding.startswith("http://") or encoding.startswith("https://"):
-
         headers = {}
         response = requests.get(encoding, timeout=30, headers=headers)
         try:

apps/shark_studio/web/index.py

@@ -1,4 +1,5 @@
 from multiprocessing import Process, freeze_support
+
 freeze_support()
 from PIL import Image
 

apps/shark_studio/web/ui/sd.py

@@ -632,7 +632,9 @@ def base_model_changed(base_model_id):
                     with gr.Tab(label="Config", id=102) as sd_tab_config:
                         with gr.Column(elem_classes=["sd-right-panel"]):
                             with gr.Row(elem_classes=["fill"]):
-                                Path(get_configs_path()).mkdir(parents=True, exist_ok=True)
+                                Path(get_configs_path()).mkdir(
+                                    parents=True, exist_ok=True
+                                )
                                 default_config_file = os.path.join(
                                     get_configs_path(),
                                     "default_sd_config.json",

apps/shark_studio/web/utils/file_utils.py

@@ -40,10 +40,12 @@
   "embeddings": {}
 }"""
 
+
 def write_default_sd_config(path):
     with open(path, "w") as f:
         f.write(default_sd_config)
 
+
 def safe_name(name):
     return name.replace("/", "_").replace("-", "_")
 

shark/iree_utils/compile_utils.py

@@ -113,8 +113,8 @@ def get_iree_frontend_args(frontend):
 # Common args to be used given any frontend or device.
 def get_iree_common_args(debug=False):
     common_args = [
-        "--iree-vm-bytecode-module-strip-source-map=true",
         "--iree-util-zero-fill-elided-attrs",
+        "--mlir-elide-elementsattrs-if-larger=10",
     ]
     if debug == True:
         common_args.extend(