Update `load_image_embeddings`

app.py

@@ -8,6 +8,7 @@ from tokenizers import Tokenizer
 from torch.utils.data import Dataset
 import albumentations as A
 from tqdm import tqdm
+from huggingface_hub import hf_hub_download
 from datasets import load_dataset
 from fourm.vq.vqvae import VQVAE
 from fourm.models.fm import FM
@@ -80,14 +81,17 @@ class HuggingFaceImageDataset(Dataset):
 # Usage
 dataset = HuggingFaceImageDataset("aroraaman/4m-21-demo")
 
-@torch.no_grad()
-def get_image_embeddings(dataset):
-    cache_file = "image_emb.pt"
-    if os.path.exists(cache_file):
-        return torch.load(cache_file)
-
-
-image_embeddings = get_image_embeddings(dataset).to(DEVICE)
+def load_image_embeddings():
+    # Download the file
+    file_path = hf_hub_download(repo_id="aroraaman/img-tensor", filename="image_emb.pt")
+    
+    # Load the tensor
+    image_embeddings = torch.load(file_path)
+    return image_embeddings
+
+# Use the embeddings in your app
+image_embeddings = load_image_embeddings()
+image_embeddings.shape
 print(image_embeddings.shape)
 
 def get_similar_images(caption, brightness, num_items):