From 037e0dbd93b5250603f533ec780af9466be234a0 Mon Sep 17 00:00:00 2001
From: Alex Nichol <unixpickle@gmail.com>
Date: Fri, 5 May 2023 12:58:16 -0700
Subject: [PATCH] checksums

---
 shap_e/models/download.py | 37 +++++++++++++++++++++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/shap_e/models/download.py b/shap_e/models/download.py
index 92c21a4..f21f7cf 100644
--- a/shap_e/models/download.py
+++ b/shap_e/models/download.py
@@ -2,6 +2,7 @@
 Adapted from: https://github.com/openai/glide-text2im/blob/69b530740eb6cef69442d6180579ef5ba9ef063e/glide_text2im/download.py
 """
 
+import hashlib
 import os
 from functools import lru_cache
 from typing import Dict, Optional
@@ -27,6 +28,18 @@ CONFIG_PATHS = {
     "diffusion": "https://openaipublic.azureedge.net/main/shap-e/diffusion_config.yaml",
 }
 
+URL_HASHES = {
+    "https://openaipublic.azureedge.net/main/shap-e/transmitter.pt": "af02a0b85a8abdfb3919584b63c540ba175f6ad4790f574a7fef4617e5acdc3b",
+    "https://openaipublic.azureedge.net/main/shap-e/vector_decoder.pt": "d7e7ebbfe3780499ae89b2da5e7c1354012dba5a6abfe295bed42f25c3be1b98",
+    "https://openaipublic.azureedge.net/main/shap-e/text_cond.pt": "e6b4fa599a7b3c3b16c222d5f5fe56f9db9289ff0b6575fbe5c11bc97106aad4",
+    "https://openaipublic.azureedge.net/main/shap-e/image_cond.pt": "cb8072c64bbbcf6910488814d212227de5db291780d4ea99c6152f9346cf12aa",
+    "https://openaipublic.azureedge.net/main/shap-e/transmitter_config.yaml": "ffe1bcb405104a37d9408391182ab118a4ef313c391e07689684f1f62071605e",
+    "https://openaipublic.azureedge.net/main/shap-e/vector_decoder_config.yaml": "e6d373649f8e24d85925f4674b9ac41c57aba5f60e42cde6d10f87381326365c",
+    "https://openaipublic.azureedge.net/main/shap-e/text_cond_config.yaml": "f290beeea3d3e9ff15db01bde5382b6e549e463060c0744f89c049505be246c1",
+    "https://openaipublic.azureedge.net/main/shap-e/image_cond_config.yaml": "4e0745605a533c543c72add803a78d233e2a6401e0abfa0cad58afb4d74ad0b0",
+    "https://openaipublic.azureedge.net/main/shap-e/diffusion_config.yaml": "efcb2cd7ee545b2d27223979d41857802448143990572a42645cd09c2942ed57",
+}
+
 
 @lru_cache()
 def default_cache_dir() -> str:
@@ -41,11 +54,14 @@ def fetch_file_cached(
     If cache_dir is specified, it will be used to download the files.
     Otherwise, default_cache_dir() is used.
     """
+    expected_hash = URL_HASHES[url]
+
     if cache_dir is None:
         cache_dir = default_cache_dir()
     os.makedirs(cache_dir, exist_ok=True)
     local_path = os.path.join(cache_dir, url.split("/")[-1])
     if os.path.exists(local_path):
+        check_hash(local_path, expected_hash)
         return local_path
 
     response = requests.get(url, stream=True)
@@ -62,9 +78,30 @@ def fetch_file_cached(
         os.rename(tmp_path, local_path)
         if progress:
             pbar.close()
+        check_hash(local_path, expected_hash)
         return local_path
 
 
+def check_hash(path: str, expected_hash: str):
+    actual_hash = hash_file(path)
+    if actual_hash != expected_hash:
+        raise RuntimeError(
+            f"The file {path} should have hash {expected_hash} but has {actual_hash}. "
+            "Try deleting it and running this call again."
+        )
+
+
+def hash_file(path: str) -> str:
+    sha256_hash = hashlib.sha256()
+    with open(path, "rb") as file:
+        while True:
+            data = file.read(4096)
+            if not len(data):
+                break
+            sha256_hash.update(data)
+    return sha256_hash.hexdigest()
+
+
 def load_config(
     config_name: str,
     progress: bool = False,