[WEB] Add vulkan-heap-block-size flag

Signed-Off-by: Gaurav Shukla <[email protected]>

Author: Gaurav Shukla

web/models/stable_diffusion/arguments.py

@@ -1,6 +1,8 @@
 from transformers import CLIPTokenizer
 from diffusers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler
 from models.stable_diffusion.opt_params import get_unet, get_vae, get_clip
+from models.stable_diffusion.utils import set_iree_runtime_flags
+import os
 
 
 class Arguments:
@@ -24,6 +26,7 @@ def __init__(
         import_mlir: bool = False,
         max_length: int = 77,
         use_tuned: bool = False,
+        vulkan_large_heap_block_size: int = 4294967296,
     ):
         self.prompt = prompt
         self.scheduler = scheduler
@@ -43,6 +46,7 @@ def __init__(
         self.import_mlir = import_mlir
         self.max_length = max_length
         self.use_tuned = use_tuned
+        self.vulkan_large_heap_block_size = vulkan_large_heap_block_size
 
     def set_params(
         self,
@@ -81,6 +85,9 @@ def set_params(
         self.import_mlir = import_mlir
 
 
+output_dir = "./stored_results/stable_diffusion"
+os.makedirs(output_dir, exist_ok=True)
+
 schedulers = dict()
 # set scheduler value
 schedulers["PNDM"] = PNDMScheduler(
@@ -112,6 +119,7 @@ def set_params(
 # cache vae and unet.
 args = Arguments()
 args.device = "vulkan"
+set_iree_runtime_flags(args)
 cache_obj["vae_fp16_vulkan"], cache_obj["unet_fp16_vulkan"] = get_vae(
     args
 ), get_unet(args)

web/models/stable_diffusion/main.py

@@ -10,10 +10,12 @@
 from numpy import iinfo
 from random import randint
 from models.stable_diffusion.opt_params import get_unet, get_vae, get_clip
-from models.stable_diffusion.arguments import args, schedulers, cache_obj
-
-output_dir = "./stored_results/stable_diffusion"
-os.makedirs(output_dir, exist_ok=True)
+from models.stable_diffusion.arguments import (
+    args,
+    schedulers,
+    cache_obj,
+    output_dir,
+)
 
 
 def stable_diff_inf(
@@ -153,6 +155,7 @@ def stable_diff_inf(
         avg_ms += step_time
         step_ms = int((step_time) * 1000)
         text_output += f"Time = {step_ms}ms."
+        print(f" \nIteration = {i}, Time = {step_ms}ms")
         latents = scheduler_obj.step(noise_pred, i, latents)["prev_sample"]
 
         if live_preview and i % 5 == 0:
@@ -178,6 +181,7 @@ def stable_diff_inf(
 
     avg_ms = 1000 * avg_ms / args.steps
     text_output += f"\n\nAverage step time: {avg_ms}ms/it"
+    print(f"\n\nAverage step time: {avg_ms}ms/it")
 
     total_time = time.time() - start
     text_output += f"\n\nTotal image generation time: {total_time}sec"

web/models/stable_diffusion/utils.py

@@ -3,6 +3,16 @@
 import torch
 from shark.shark_inference import SharkInference
 from shark.shark_importer import import_with_fx
+from shark.iree_utils.vulkan_utils import set_iree_vulkan_runtime_flags
+
+
+def set_iree_runtime_flags(args):
+    vulkan_runtime_flags = [
+        f"--vulkan_large_heap_block_size={args.vulkan_large_heap_block_size}",
+    ]
+    if "vulkan" in args.device:
+        set_iree_vulkan_runtime_flags(flags=vulkan_runtime_flags)
+    return
 
 
 def _compile_module(args, shark_module, model_name, extra_args=[]):