diff --git a/model_filemanager/download_models.py b/model_filemanager/download_models.py index 712d59328..b24cbba59 100644 --- a/model_filemanager/download_models.py +++ b/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_url: str, + model_name: 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): - The subdirectory within the main models directory where the model + model_sub_directory (str): + 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, - progress_callback: Callable[[str, DownloadModelStatus], Awaitable[Any]], + +async def check_file_exists(file_path: str, + 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, - file_path: str, - model_name: str, - progress_callback: Callable[[str, DownloadModelStatus], Awaitable[Any]], - relative_path: str, +async def track_download_progress(response: aiohttp.ClientResponse, + file_path: str, + model_name: str, + progress_callback: Callable[[str, DownloadModelStatus], Awaitable[Any]], + 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, - progress_callback: Callable[[str, DownloadModelStatus], Any], + +async def handle_download_error(e: Exception, + 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