(minor) format download_models.py

model_filemanager/download_models.py

@@ -17,6 +17,7 @@ class DownloadStatusType(Enum):
     COMPLETED = "completed"
     ERROR = "error"
 
+
 @dataclass
 class DownloadModelStatus():
     status: str
@@ -29,7 +30,7 @@ class DownloadModelStatus():
         self.progress_percentage = progress_percentage
         self.message = message
         self.already_existed = already_existed
-    
+
     def to_dict(self) -> Dict[str, Any]:
         return {
             "status": self.status,
@@ -38,9 +39,10 @@ class DownloadModelStatus():
             "already_existed": self.already_existed
         }
 
+
 async def download_model(model_download_request: Callable[[str], Awaitable[aiohttp.ClientResponse]],
-                         model_name: str,  
+                         model_name: str,
-                         model_url: str, 
+                         model_url: str,
                          model_sub_directory: str,
                          progress_callback: Callable[[str, DownloadModelStatus], Awaitable[Any]],
                          progress_interval: float = 1.0) -> DownloadModelStatus:
@@ -48,16 +50,16 @@ async def download_model(model_download_request: Callable[[str], Awaitable[aioht
     Download a model file from a given URL into the models directory.
 
     Args:
-        model_download_request (Callable[[str], Awaitable[aiohttp.ClientResponse]]): 
+        model_download_request (Callable[[str], Awaitable[aiohttp.ClientResponse]]):
             A function that makes an HTTP request. This makes it easier to mock in unit tests.
-        model_name (str): 
+        model_name (str):
             The name of the model file to be downloaded. This will be the filename on disk.
-        model_url (str): 
+        model_url (str):
             The URL from which to download the model.
-        model_sub_directory (str): 
+        model_sub_directory (str):
-            The subdirectory within the main models directory where the model 
+            The subdirectory within the main models directory where the model
             should be saved (e.g., 'checkpoints', 'loras', etc.).
-        progress_callback (Callable[[str, DownloadModelStatus], Awaitable[Any]]): 
+        progress_callback (Callable[[str, DownloadModelStatus], Awaitable[Any]]):
             An asynchronous function to call with progress updates.
 
     Returns:
@@ -65,17 +67,17 @@ async def download_model(model_download_request: Callable[[str], Awaitable[aioht
     """
     if not validate_model_subdirectory(model_sub_directory):
         return DownloadModelStatus(
-            DownloadStatusType.ERROR, 
+            DownloadStatusType.ERROR,
             0,
-            "Invalid model subdirectory", 
+            "Invalid model subdirectory",
             False
         )
 
     if not validate_filename(model_name):
         return DownloadModelStatus(
-            DownloadStatusType.ERROR, 
+            DownloadStatusType.ERROR,
             0,
-            "Invalid model name", 
+            "Invalid model name",
             False
         )
 
@@ -101,7 +103,7 @@ async def download_model(model_download_request: Callable[[str], Awaitable[aioht
     except Exception as e:
         logging.error(f"Error in downloading model: {e}")
         return await handle_download_error(e, model_name, progress_callback, relative_path)
-    
+
 
 def create_model_path(model_name: str, model_directory: str, models_base_dir: str) -> tuple[str, str]:
     full_model_dir = os.path.join(models_base_dir, model_directory)
@@ -114,13 +116,13 @@ def create_model_path(model_name: str, model_directory: str, models_base_dir: st
     if os.path.commonprefix([abs_file_path, abs_base_dir]) != abs_base_dir:
         raise Exception(f"Invalid model directory: {model_directory}/{model_name}")
 
-
     relative_path = '/'.join([model_directory, model_name])
     return file_path, relative_path
 
-async def check_file_exists(file_path: str, 
+
-                            model_name: str, 
+async def check_file_exists(file_path: str,
-                            progress_callback: Callable[[str, DownloadModelStatus], Awaitable[Any]], 
+                            model_name: str,
+                            progress_callback: Callable[[str, DownloadModelStatus], Awaitable[Any]],
                             relative_path: str) -> Optional[DownloadModelStatus]:
     if os.path.exists(file_path):
         status = DownloadModelStatus(DownloadStatusType.COMPLETED, 100, f"{model_name} already exists", True)
@@ -129,11 +131,11 @@ async def check_file_exists(file_path: str,
     return None
 
 
-async def track_download_progress(response: aiohttp.ClientResponse, 
+async def track_download_progress(response: aiohttp.ClientResponse,
-                                  file_path: str, 
+                                  file_path: str,
-                                  model_name: str, 
+                                  model_name: str,
-                                  progress_callback: Callable[[str, DownloadModelStatus], Awaitable[Any]], 
+                                  progress_callback: Callable[[str, DownloadModelStatus], Awaitable[Any]],
-                                  relative_path: str, 
+                                  relative_path: str,
                                   interval: float = 1.0) -> DownloadModelStatus:
     try:
         total_size = int(response.headers.get('Content-Length', 0))
@@ -156,12 +158,12 @@ async def track_download_progress(response: aiohttp.ClientResponse,
                     break
                 f.write(chunk)
                 downloaded += len(chunk)
-                
+
                 if time.time() - last_update_time >= interval:
                     await update_progress()
 
         await update_progress()
-        
+
         logging.info(f"Successfully downloaded {model_name}. Total downloaded: {downloaded}")
         status = DownloadModelStatus(DownloadStatusType.COMPLETED, 100, f"Successfully downloaded {model_name}", False)
         await progress_callback(relative_path, status)
@@ -172,18 +174,20 @@ async def track_download_progress(response: aiohttp.ClientResponse,
         logging.error(traceback.format_exc())
         return await handle_download_error(e, model_name, progress_callback, relative_path)
 
-async def handle_download_error(e: Exception, 
+
-                                model_name: str, 
+async def handle_download_error(e: Exception,
-                                progress_callback: Callable[[str, DownloadModelStatus], Any], 
+                                model_name: str,
+                                progress_callback: Callable[[str, DownloadModelStatus], Any],
                                 relative_path: str) -> DownloadModelStatus:
     error_message = f"Error downloading {model_name}: {str(e)}"
     status = DownloadModelStatus(DownloadStatusType.ERROR, 0, error_message, False)
     await progress_callback(relative_path, status)
     return status
 
+
 def validate_model_subdirectory(model_subdirectory: str) -> bool:
     """
-    Validate that the model subdirectory is safe to install into. 
+    Validate that the model subdirectory is safe to install into.
     Must not contain relative paths, nested paths or special characters
     other than underscores and hyphens.
 
@@ -204,10 +208,11 @@ def validate_model_subdirectory(model_subdirectory: str) -> bool:
 
     return True
 
+
 def validate_filename(filename: str)-> bool:
     """
     Validate a filename to ensure it's safe and doesn't contain any path traversal attempts.
-    
+
     Args:
     filename (str): The filename to validate