Revert "find gsutil on linux (huggingface#557)" (huggingface#560)

This reverts commit bba8646.

Author: dan-garvey

cpp/save_img.py

@@ -1,6 +1,7 @@
 import numpy as np
 import tensorflow as tf
 from shark.shark_inference import SharkInference
+from shark.shark_downloader import download_tf_model
 
 
 def load_and_preprocess_image(fname: str):

requirements.txt

@@ -6,7 +6,7 @@ pyinstaller
 tqdm
 
 # SHARK Downloader
-google-cloud-storage
+gsutil
 
 # Testing
 pytest

shark/examples/shark_inference/bloom_tank.py

@@ -1,9 +1,7 @@
 from shark.shark_inference import SharkInference
-from shark.shark_downloader import download_model
+from shark.shark_downloader import download_torch_model
 
-mlir_model, func_name, inputs, golden_out = download_model(
-    "bloom", frontend="torch"
-)
+mlir_model, func_name, inputs, golden_out = download_torch_model("bloom")
 
 shark_module = SharkInference(
     mlir_model, func_name, device="cpu", mlir_dialect="tm_tensor"

shark/examples/shark_inference/minilm_jit.py

@@ -1,10 +1,9 @@
 from shark.shark_inference import SharkInference
-from shark.shark_downloader import download_model
+from shark.shark_downloader import download_torch_model
 
 
-mlir_model, func_name, inputs, golden_out = download_model(
-    "microsoft/MiniLM-L12-H384-uncased",
-    frontend="torch",
+mlir_model, func_name, inputs, golden_out = download_torch_model(
+    "microsoft/MiniLM-L12-H384-uncased"
 )
 
 

shark/examples/shark_inference/resnet50_script.py

@@ -5,7 +5,7 @@
 from torchvision import transforms
 import sys
 from shark.shark_inference import SharkInference
-from shark.shark_downloader import download_model
+from shark.shark_downloader import download_torch_model
 
 
 ################################## Preprocessing inputs and model ############
@@ -66,9 +66,7 @@ def forward(self, img):
 
 
 ## Can pass any img or input to the forward module.
-mlir_model, func_name, inputs, golden_out = download_model(
-    "resnet50", frontend="torch"
-)
+mlir_model, func_name, inputs, golden_out = download_torch_model("resnet50")
 
 shark_module = SharkInference(mlir_model, func_name, mlir_dialect="linalg")
 shark_module.compile()

shark/examples/shark_inference/stable_diff_f16.py

@@ -37,12 +37,10 @@
 
 
 def fp16_unet():
-    from shark.shark_downloader import download_model
+    from shark.shark_downloader import download_torch_model
 
-    mlir_model, func_name, inputs, golden_out = download_model(
-        "stable_diff_f16_18_OCT",
-        tank_url="gs://shark_tank/prashant_nod",
-        frontend="torch",
+    mlir_model, func_name, inputs, golden_out = download_torch_model(
+        "stable_diff_f16_18_OCT", tank_url="gs://shark_tank/prashant_nod"
     )
     shark_module = SharkInference(
         mlir_model, func_name, device=args.device, mlir_dialect="linalg"

shark/examples/shark_inference/stable_diff_tf.py

@@ -17,7 +17,7 @@
 )
 
 from shark.shark_inference import SharkInference
-from shark.shark_downloader import download_model
+from shark.shark_downloader import download_tf_model
 from PIL import Image
 
 # pip install "git+https://github.com/keras-team/keras-cv.git"
@@ -75,8 +75,8 @@ def __init__(self, device="cpu", jit_compile=True):
         # Create models
         self.text_encoder = TextEncoder(MAX_PROMPT_LENGTH)
 
-        mlir_model, func_name, inputs, golden_out = download_model(
-            "stable_diff", tank_url="gs://shark_tank/quinn", frontend="tf"
+        mlir_model, func_name, inputs, golden_out = download_tf_model(
+            "stable_diff", tank_url="gs://shark_tank/quinn"
         )
         shark_module = SharkInference(
             mlir_model, func_name, device=device, mlir_dialect="mhlo"

shark/examples/shark_inference/stable_diffusion/utils.py

@@ -39,12 +39,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.shark_downloader import download_torch_model
 
-    mlir_model, func_name, inputs, golden_out = download_model(
-        model_name,
-        tank_url=tank_url,
-        frontend="torch",
+    mlir_model, func_name, inputs, golden_out = download_torch_model(
+        model_name, tank_url=tank_url
     )
     shark_module = SharkInference(
         mlir_model, func_name, device=args.device, mlir_dialect="linalg"

shark/examples/shark_inference/v_diffusion.py

@@ -1,10 +1,8 @@
 from shark.shark_inference import SharkInference
-from shark.shark_downloader import download_model
+from shark.shark_downloader import download_torch_model
 
 
-mlir_model, func_name, inputs, golden_out = download_model(
-    "v_diffusion", frontend="torch"
-)
+mlir_model, func_name, inputs, golden_out = download_torch_model("v_diffusion")
 
 shark_module = SharkInference(
     mlir_model, func_name, device="vulkan", mlir_dialect="linalg"

shark/shark_downloader.py

@@ -17,22 +17,17 @@
 import sys
 from pathlib import Path
 from shark.parser import shark_args
-from google.cloud import storage
 
 
-def download_public_file(full_gs_url, destination_file_name):
-    """Downloads a public blob from the bucket."""
-    # bucket_name = "gs://your-bucket-name/path/to/file"
-    # destination_file_name = "local/path/to/file"
-
-    storage_client = storage.Client.create_anonymous_client()
-    bucket_name = full_gs_url.split("/")[2]
-    source_blob_name = "/".join(full_gs_url.split("/")[3:])
-    bucket = storage_client.bucket(bucket_name)
-    blob = bucket.blob(source_blob_name)
-    blob.download_to_filename(destination_file_name)
+def resource_path(relative_path):
+    """Get absolute path to resource, works for dev and for PyInstaller"""
+    base_path = getattr(
+        sys, "_MEIPASS", os.path.dirname(os.path.abspath(__file__))
+    )
+    return os.path.join(base_path, relative_path)
 
 
+GSUTIL_PATH = resource_path("gsutil")
 GSUTIL_FLAGS = ' -o "GSUtil:parallel_process_count=1" -m cp -r '
 
 
@@ -103,23 +98,103 @@ def check_dir_exists(model_name, frontend="torch", dynamic=""):
 
 
 # Downloads the torch model from gs://shark_tank dir.
-def download_model(
-    model_name,
-    dynamic=False,
-    tank_url="gs://shark_tank/latest",
-    frontend=None,
-    tuned=None,
+def download_torch_model(
+    model_name, dynamic=False, tank_url="gs://shark_tank/latest"
 ):
     model_name = model_name.replace("/", "_")
     dyn_str = "_dynamic" if dynamic else ""
     os.makedirs(WORKDIR, exist_ok=True)
-    model_dir_name = model_name + "_" + frontend
-    full_gs_url = tank_url.rstrip("/") + "/" + model_dir_name
+    model_dir_name = model_name + "_torch"
+
+    def gs_download_model():
+        gs_command = (
+            GSUTIL_PATH
+            + GSUTIL_FLAGS
+            + tank_url
+            + "/"
+            + model_dir_name
+            + ' "'
+            + WORKDIR
+            + '"'
+        )
+        if os.system(gs_command) != 0:
+            raise Exception("model not present in the tank. Contact Nod Admin")
+
+    if not check_dir_exists(model_dir_name, frontend="torch", dynamic=dyn_str):
+        gs_download_model()
+    else:
+        if not _internet_connected():
+            print(
+                "No internet connection. Using the model already present in the tank."
+            )
+        else:
+            model_dir = os.path.join(WORKDIR, model_dir_name)
+            local_hash = str(np.load(os.path.join(model_dir, "hash.npy")))
+            gs_hash = (
+                GSUTIL_PATH
+                + GSUTIL_FLAGS
+                + tank_url
+                + "/"
+                + model_dir_name
+                + "/hash.npy"
+                + " "
+                + os.path.join(model_dir, "upstream_hash.npy")
+            )
+            if os.system(gs_hash) != 0:
+                raise Exception("hash of the model not present in the tank.")
+            upstream_hash = str(
+                np.load(os.path.join(model_dir, "upstream_hash.npy"))
+            )
+            if local_hash != upstream_hash:
+                if shark_args.update_tank == True:
+                    gs_download_model()
+                else:
+                    print(
+                        "Hash does not match upstream in gs://shark_tank/. If you are using SHARK Downloader with locally generated artifacts, this is working as intended."
+                    )
+
+    model_dir = os.path.join(WORKDIR, model_dir_name)
+    with open(
+        os.path.join(model_dir, model_name + dyn_str + "_torch.mlir"),
+        mode="rb",
+    ) as f:
+        mlir_file = f.read()
+
+    function_name = str(np.load(os.path.join(model_dir, "function_name.npy")))
+    inputs = np.load(os.path.join(model_dir, "inputs.npz"))
+    golden_out = np.load(os.path.join(model_dir, "golden_out.npz"))
+
+    inputs_tuple = tuple([inputs[key] for key in inputs])
+    golden_out_tuple = tuple([golden_out[key] for key in golden_out])
+    return mlir_file, function_name, inputs_tuple, golden_out_tuple
+
+
+# Downloads the tflite model from gs://shark_tank dir.
+def download_tflite_model(
+    model_name, dynamic=False, tank_url="gs://shark_tank/latest"
+):
+    dyn_str = "_dynamic" if dynamic else ""
+    os.makedirs(WORKDIR, exist_ok=True)
+    model_dir_name = model_name + "_tflite"
+
+    def gs_download_model():
+        gs_command = (
+            GSUTIL_PATH
+            + GSUTIL_FLAGS
+            + tank_url
+            + "/"
+            + model_dir_name
+            + ' "'
+            + WORKDIR
+            + '"'
+        )
+        if os.system(gs_command) != 0:
+            raise Exception("model not present in the tank. Contact Nod Admin")
 
     if not check_dir_exists(
-        model_dir_name, frontend=frontend, dynamic=dyn_str
+        model_dir_name, frontend="tflite", dynamic=dyn_str
     ):
-        download_public_file(full_gs_url, WORKDIR)
+        gs_download_model()
     else:
         if not _internet_connected():
             print(
@@ -128,34 +203,104 @@ def download_model(
         else:
             model_dir = os.path.join(WORKDIR, model_dir_name)
             local_hash = str(np.load(os.path.join(model_dir, "hash.npy")))
-            gs_hash_url = (
-                tank_url.rstrip("/") + "/" + model_dir_name + "/hash.npy"
+            gs_hash = (
+                GSUTIL_PATH
+                + GSUTIL_FLAGS
+                + tank_url
+                + "/"
+                + model_dir_name
+                + "/hash.npy"
+                + " "
+                + os.path.join(model_dir, "upstream_hash.npy")
             )
-            download_public_file(
-                gs_hash_url, os.path.join(model_dir, "upstream_hash.npy")
+            if os.system(gs_hash) != 0:
+                raise Exception("hash of the model not present in the tank.")
+            upstream_hash = str(
+                np.load(os.path.join(model_dir, "upstream_hash.npy"))
             )
+            if local_hash != upstream_hash:
+                if shark_args.update_tank == True:
+                    gs_download_model()
+                else:
+                    print(
+                        "Hash does not match upstream in gs://shark_tank/. If you are using SHARK Downloader with locally generated artifacts, this is working as intended."
+                    )
+
+    model_dir = os.path.join(WORKDIR, model_dir_name)
+    with open(
+        os.path.join(model_dir, model_name + dyn_str + "_tflite.mlir"),
+        mode="rb",
+    ) as f:
+        mlir_file = f.read()
+
+    function_name = str(np.load(os.path.join(model_dir, "function_name.npy")))
+    inputs = np.load(os.path.join(model_dir, "inputs.npz"))
+    golden_out = np.load(os.path.join(model_dir, "golden_out.npz"))
+
+    inputs_tuple = tuple([inputs[key] for key in inputs])
+    golden_out_tuple = tuple([golden_out[key] for key in golden_out])
+    return mlir_file, function_name, inputs_tuple, golden_out_tuple
+
+
+def download_tf_model(
+    model_name, tuned=None, tank_url="gs://shark_tank/latest"
+):
+    model_name = model_name.replace("/", "_")
+    os.makedirs(WORKDIR, exist_ok=True)
+    model_dir_name = model_name + "_tf"
+
+    def gs_download_model():
+        gs_command = (
+            GSUTIL_PATH
+            + GSUTIL_FLAGS
+            + tank_url
+            + "/"
+            + model_dir_name
+            + ' "'
+            + WORKDIR
+            + '"'
+        )
+        if os.system(gs_command) != 0:
+            raise Exception("model not present in the tank. Contact Nod Admin")
+
+    if not check_dir_exists(model_dir_name, frontend="tf"):
+        gs_download_model()
+    else:
+        if not _internet_connected():
+            print(
+                "No internet connection. Using the model already present in the tank."
+            )
+        else:
+            model_dir = os.path.join(WORKDIR, model_dir_name)
+            local_hash = str(np.load(os.path.join(model_dir, "hash.npy")))
+            gs_hash = (
+                GSUTIL_PATH
+                + GSUTIL_FLAGS
+                + tank_url
+                + "/"
+                + model_dir_name
+                + "/hash.npy"
+                + " "
+                + os.path.join(model_dir, "upstream_hash.npy")
+            )
+            if os.system(gs_hash) != 0:
+                raise Exception("hash of the model not present in the tank.")
             upstream_hash = str(
                 np.load(os.path.join(model_dir, "upstream_hash.npy"))
             )
             if local_hash != upstream_hash:
                 if shark_args.update_tank == True:
-                    download_public_file(full_gs_url, WORKDIR)
+                    gs_download_model()
                 else:
                     print(
                         "Hash does not match upstream in gs://shark_tank/. If you are using SHARK Downloader with locally generated artifacts, this is working as intended."
                     )
 
     model_dir = os.path.join(WORKDIR, model_dir_name)
-    suffix = (
-        "_" + frontend + ".mlir"
-        if tuned is None
-        else "_" + frontend + "_" + tuned + ".mlir"
-    )
+    suffix = "_tf.mlir" if tuned is None else "_tf_" + tuned + ".mlir"
     filename = os.path.join(model_dir, model_name + suffix)
     if not os.path.isfile(filename):
-        filename = os.path.join(
-            model_dir, model_name + "_" + frontend + ".mlir"
-        )
+        filename = os.path.join(model_dir, model_name + "_tf.mlir")
 
     with open(filename, mode="rb") as f:
         mlir_file = f.read()