feat: return model_paths in fetch_models_task

Signed-off-by: Anupam Kumar <[email protected]>

Author: kyteinsky

nc_py_api/ex_app/integration_fastapi.py

@@ -120,29 +120,35 @@ def __map_app_static_folders(fast_api_app: FastAPI):
             fast_api_app.mount(f"/{mnt_dir}", staticfiles.StaticFiles(directory=mnt_dir_path), name=mnt_dir)
 
 
-def fetch_models_task(nc: NextcloudApp, models: dict[str, dict], progress_init_start_value: int) -> None:
+def fetch_models_task(
+    nc: NextcloudApp, models: dict[str, dict], progress_init_start_value: int
+) -> dict[str, str | None]:
     """Use for cases when you want to define custom `/init` but still need to easy download models."""
+    model_paths = {}
     if models:
         current_progress = progress_init_start_value
         percent_for_each = min(int((100 - progress_init_start_value) / len(models)), 99)
         for model in models:
             if model.startswith(("http://", "https://")):
-                __fetch_model_as_file(current_progress, percent_for_each, nc, model, models[model])
+                model_paths[model] = __fetch_model_as_file(current_progress, percent_for_each, nc, model, models[model])
             else:
-                __fetch_model_as_snapshot(current_progress, percent_for_each, nc, model, models[model])
+                model_paths[model] = __fetch_model_as_snapshot(
+                    current_progress, percent_for_each, nc, model, models[model]
+                )
             current_progress += percent_for_each
     nc.set_init_status(100)
+    return model_paths
 
 
 def __fetch_model_as_file(
     current_progress: int, progress_for_task: int, nc: NextcloudApp, model_path: str, download_options: dict
-) -> None:
+) -> str | None:
     result_path = download_options.pop("save_path", urlparse(model_path).path.split("/")[-1])
     try:
         with httpx.stream("GET", model_path, follow_redirects=True) as response:
             if not response.is_success:
                 nc.log(LogLvl.ERROR, f"Downloading of '{model_path}' returned {response.status_code} status.")
-                return
+                return None
             downloaded_size = 0
             linked_etag = ""
             for each_history in response.history:
@@ -163,7 +169,7 @@ def __fetch_model_as_file(
                         sha256_hash.update(byte_block)
                     if f'"{sha256_hash.hexdigest()}"' == linked_etag:
                         nc.set_init_status(min(current_progress + progress_for_task, 99))
-                        return
+                        return None
 
             with builtins.open(result_path, "wb") as file:
                 last_progress = current_progress
@@ -174,13 +180,17 @@ def __fetch_model_as_file(
                         if new_progress != last_progress:
                             nc.set_init_status(new_progress)
                             last_progress = new_progress
+
+        return result_path
     except Exception as e:  # noqa pylint: disable=broad-exception-caught
         nc.log(LogLvl.ERROR, f"Downloading of '{model_path}' raised an exception: {e}")
 
+    return None
+
 
 def __fetch_model_as_snapshot(
     current_progress: int, progress_for_task, nc: NextcloudApp, mode_name: str, download_options: dict
-) -> None:
+) -> str:
     from huggingface_hub import snapshot_download  # noqa isort:skip pylint: disable=C0415 disable=E0401
     from tqdm import tqdm  # noqa isort:skip pylint: disable=C0415 disable=E0401
 
@@ -191,7 +201,9 @@ def display(self, msg=None, pos=None):
 
     workers = download_options.pop("max_workers", 2)
     cache = download_options.pop("cache_dir", persistent_storage())
-    snapshot_download(mode_name, tqdm_class=TqdmProgress, **download_options, max_workers=workers, cache_dir=cache)
+    return snapshot_download(
+        mode_name, tqdm_class=TqdmProgress, **download_options, max_workers=workers, cache_dir=cache
+    )
 
 
 def __nc_app(request: HTTPConnection) -> dict: