Update app.py

app.py

@@ -23,7 +23,9 @@ from typing import Dict, List, Optional
 from huggingface_hub import login, HfApi, hf_hub_download  # Import hf_hub_download
 from huggingface_hub.utils import validate_repo_id, HFValidationError
 from huggingface_hub.errors import HfHubHTTPError
-
+from huggingface_hub import HfApi, hf_hub_download, cached_download, get_from_cache  # Import cached_download and get_from_cache
+from huggingface_hub.utils import validate_repo_id, HFValidationError
+from huggingface_hub.constants import HUGGINGFACE_HUB_CACHE
 
 # ---------------------- DEPENDENCIES ----------------------
 def install_dependencies_gradio():
@@ -58,26 +60,50 @@ def create_model_repo(api, user, orgs_name, model_name, make_private=False):
 
 # ---------------------- MODEL LOADING AND CONVERSION ----------------------
 def download_model(model_path_or_url):
-    """Downloads a model, handling URLs, HF repos, and local paths."""
+    """Downloads a model, handling URLs, HF repos, and local paths, caching appropriately."""
     try:
-        # Check if it's a valid Hugging Face repo ID (and potentially a file within)
+        # 1. Check if it's a valid Hugging Face repo ID (and potentially a file within)
         try:
             validate_repo_id(model_path_or_url)
-            # It's a valid repo ID; use hf_hub_download without a filename
+            # It's a valid repo ID; use hf_hub_download
             local_path = hf_hub_download(repo_id=model_path_or_url)
             return local_path
         except HFValidationError:
-            pass  # Not a simple repo ID.  Might be repo ID + filename, or a URL.
+            pass  # Not a simple repo ID. Might be repo ID + filename, or a URL.
 
+        # 2. Check if it's a URL
         if model_path_or_url.startswith("http://") or model_path_or_url.startswith("https://"):
-            # It's a URL: use hf_hub_download with the url parameter
-            local_path = hf_hub_download(url=model_path_or_url)  # Corrected line
+            # Check if it's already in the cache
+            cache_path = get_from_cache(model_path_or_url)  # Use get_from_cache
+            if cache_path is not None:
+                return cache_path
+
+            # It's a URL and not in cache: download manually and put into HF cache
+            response = requests.get(model_path_or_url, stream=True)
+            response.raise_for_status()  # Raise HTTPError for bad requests (4xx or 5xx)
+
+            # Get filename from URL, or use a hash if we can't determine it
+            parsed_url = urlparse(model_path_or_url)
+            filename = os.path.basename(unquote(parsed_url.path))
+            if not filename:
+              filename = hashlib.sha256(model_path_or_url.encode()).hexdigest()
+
+            # Construct the cache path (using HF_HUB_CACHE + "downloads" )
+            cache_dir = os.path.join(HUGGINGFACE_HUB_CACHE, "downloads")
+            os.makedirs(cache_dir, exist_ok=True) # Ensure the cache directory exists
+            local_path = os.path.join(cache_dir, filename)
+
+            with open(local_path, "wb") as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    f.write(chunk)
             return local_path
+
+        # 3. Check if it's a local file
         elif os.path.isfile(model_path_or_url):
-            # It's a local file
             return model_path_or_url
+
+        # 4. Handle Hugging Face repo with a specific file
         else:
-            # Try splitting into repo ID and filename
             try:
                 parts = model_path_or_url.split("/", 1)
                 if len(parts) == 2:
@@ -86,7 +112,8 @@ def download_model(model_path_or_url):
                     local_path = hf_hub_download(repo_id=repo_id, filename=filename)
                     return local_path
                 else:
-                    raise ValueError("Invalid input format")
+                   raise ValueError("Invalid input format.")
+
             except HFValidationError:
                 raise ValueError(f"Invalid model path or URL: {model_path_or_url}")