Update hub_models.py

hub_models.py

@@ -7,6 +7,8 @@ import io
 import json
 import os
 import re
+import shutil
+from threading import Lock
 import time
 
 CACHE_ROOT = Path(__file__).resolve().parent / '.model-cache'
@@ -17,6 +19,11 @@ from huggingface_hub.utils import validate_repo_id
 from tqdm.auto import tqdm
 
 from embeddings import embedding_name
+from app_environment import hosted_mode, hosted_download_limit
+
+
+HOSTED_DOWNLOAD_LOCK = Lock()
+DISK_RESERVE_BYTES = 2_000_000_000
 
 
 @dataclass(frozen=True)
@@ -47,6 +54,7 @@ class DownloadPlan:
     directory: Path
     size: int | None
     embeddings_only: bool = False
+    file_sizes: tuple[int | None, ...] = ()
 
 
 def safe_filename(name):
@@ -151,11 +159,54 @@ def build_plan(catalog, selection, token=None, cache_root=CACHE_ROOT):
     size = sum(sizes) if all(value is not None for value in sizes) else None
     model_path = directory if selection.endswith('.index.json') else directory / selection
     return DownloadPlan(catalog.repo_id, catalog.commit, selection, files, tuple(weights), model_path,
-                         directory / tokenizer if tokenizer else None, directory, size, embeddings_only)
+                         directory / tokenizer if tokenizer else None, directory, size, embeddings_only, tuple(sizes))
+
+
+def check_hosted_download(plan):
+    """Enforce the per-model limit and leave room for the running Space.
+
+    Complete cached files need no additional disk space. Incomplete downloads
+    are counted at their full size, making this estimate conservative.
+    """
+    limit = hosted_download_limit()
+    if limit is None:
+        return
+    if plan.size is None or plan.size > limit:
+        raise ValueError(f'This Space supports models up to {format_size(limit)} with a known download size. '
+                         'Choose smaller weights or embedding shards, or run the app locally.')
+    remaining = plan.size
+    if len(plan.file_sizes) == len(plan.files):
+        for name, size in zip(plan.files, plan.file_sizes):
+            safe_filename(name)
+            path = plan.directory / name
+            if size is not None and path.is_file() and path.stat().st_size == size:
+                remaining -= size
+    ancestor = plan.directory.resolve()
+    while not ancestor.exists():
+        ancestor = ancestor.parent
+    free = shutil.disk_usage(ancestor).free
+    if remaining and remaining + DISK_RESERVE_BYTES > free:
+        raise ValueError(f'Not enough free storage on this Space: {format_size(remaining)} still needs downloading, '
+                         f'but only {format_size(free)} is free. We reserve 2 GB for the app. '
+                         'Choose smaller weights or try again after the Space cache is cleared.')
 
 
 def download_model(plan, token=None, progress=None):
     """Progress callbacks may run on SDK threads; callers should use a queue."""
+    locked = hosted_mode()
+    if locked and not HOSTED_DOWNLOAD_LOCK.acquire(blocking=False):
+        raise ValueError('Another model is downloading on this Space. Please try again when it finishes.')
+    try:
+        # Recheck inside the lock: another visitor may have consumed disk space
+        # since this session displayed the download plan.
+        check_hosted_download(plan)
+        return _download_model(plan, token, progress)
+    finally:
+        if locked:
+            HOSTED_DOWNLOAD_LOCK.release()
+
+
+def _download_model(plan, token=None, progress=None):
     total_files = len(plan.files)
     for index, filename in enumerate(plan.files):
         safe_filename(filename)