[SHARK][SD] Add --local_tank_cache flag in the stable diffusion

This flag can be used to set local shark_tank cache directory.

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

Author: Shukla-Gaurav

shark/examples/shark_inference/stable_diffusion/stable_args.py

@@ -103,6 +103,12 @@
     help="Download and use the tuned version of the model if available",
 )
 
+p.add_argument(
+    "--local_tank_cache",
+    default="",
+    help="Specify where to save downloaded shark_tank artifacts. If this is not set, the default is ~/.local/shark_tank/.",
+)
+
 p.add_argument(
     "--dump_isa",
     default=False,

shark/examples/shark_inference/stable_diffusion/utils.py

@@ -40,6 +40,10 @@ def _compile_module(shark_module, model_name, extra_args=[]):
 # Downloads the model from shark_tank and returns the shark_module.
 def get_shark_model(tank_url, model_name, extra_args=[]):
     from shark.shark_downloader import download_model
+    from shark.parser import shark_args
+
+    # Set local shark_tank cache directory.
+    shark_args.local_tank_cache = args.local_tank_cache
 
     mlir_model, func_name, inputs, golden_out = download_model(
         model_name,

web/models/stable_diffusion/stable_args.py

@@ -110,6 +110,12 @@
     help="Download and use the tuned version of the model if available",
 )
 
+p.add_argument(
+    "--local_tank_cache",
+    default="",
+    help="Specify where to save downloaded shark_tank artifacts. If this is not set, the default is ~/.local/shark_tank/.",
+)
+
 p.add_argument(
     "--vulkan_large_heap_block_size",
     default="4294967296",

web/models/stable_diffusion/utils.py

@@ -38,6 +38,10 @@ def _compile_module(args, shark_module, model_name, extra_args=[]):
 # Downloads the model from shark_tank and returns the shark_module.
 def get_shark_model(args, tank_url, model_name, extra_args=[]):
     from shark.shark_downloader import download_model
+    from shark.parser import shark_args
+
+    # Set local shark_tank cache directory.
+    shark_args.local_tank_cache = args.local_tank_cache
 
     mlir_model, func_name, inputs, golden_out = download_model(
         model_name, tank_url=tank_url, frontend="torch"