precache clip

aria/image_encoder.py

@@ -3,6 +3,7 @@ import torch.nn as nn
 from transformers import CLIPProcessor, CLIPModel
 from PIL import Image
 from typing import Tuple, Union
+import os
 
 class ImageEncoder(nn.Module):
     def __init__(self, clip_model_name: str = "openai/clip-vit-large-patch14-336"):
@@ -30,15 +31,39 @@ class ImageEncoder(nn.Module):
         print(f"Initializing ImageEncoder with {self.clip_model_name}...")
         print("Loading CLIP model from local cache (network disabled)...")
 
-        # Load CLIP model and processor strictly from the local Hugging-Face cache
-        self.clip_model = CLIPModel.from_pretrained(
-            self.clip_model_name,
-            local_files_only=True  # fail fast if cache is missing
-        )
-        self.processor = CLIPProcessor.from_pretrained(
-            self.clip_model_name,
-            local_files_only=True
-        )
+        # Prefer loading strictly from the local Hugging Face cache that `app.py` populates.
+        # If the files are genuinely missing (e.g. first run without network), we fall back
+        # to an online download so the user still gets a working application.
+
+        # Determine the cache directory from env – this is set in `app.py`.
+        hf_cache_dir = os.environ.get("HF_HUB_CACHE", None)
+
+        try:
+            self.clip_model = CLIPModel.from_pretrained(
+                self.clip_model_name,
+                cache_dir=hf_cache_dir,
+                local_files_only=True,  # use cache only on the first attempt
+            )
+            self.processor = CLIPProcessor.from_pretrained(
+                self.clip_model_name,
+                cache_dir=hf_cache_dir,
+                local_files_only=True,
+            )
+            print("CLIP model loaded successfully from local cache")
+        except (OSError, EnvironmentError) as cache_err:
+            print(
+                "Local cache for CLIP model not found – attempting a one-time online download..."
+            )
+            # Note: this will still respect HF_HUB_CACHE so the files are cached for future runs.
+            self.clip_model = CLIPModel.from_pretrained(
+                self.clip_model_name,
+                cache_dir=hf_cache_dir,
+            )
+            self.processor = CLIPProcessor.from_pretrained(
+                self.clip_model_name,
+                cache_dir=hf_cache_dir,
+            )
+            print("CLIP model downloaded and cached successfully")
         
         print("CLIP model loaded successfully")