xinyun/ComfyUI
1
0
Fork 0
mirror of https://git.datalinker.icu/comfyanonymous/ComfyUI synced 2026-11-10 14:00:05 +08:00
Code Issues Packages Projects Releases Wiki Activity

(minor) format download_models.py

Browse Source
This commit is contained in:
Alex "mcmonkey" Goodwin 2024-09-19 19:00:16 +09:00
parent de8e8e3b0d
commit cbaffac4a5
1 changed files with 35 additions and 30 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
model_filemanager/download_models.py
View File

@ -17,6 +17,7 @@ class DownloadStatusType(Enum):
COMPLETED = "completed" COMPLETED = "completed"
ERROR = "error" ERROR = "error"
@dataclass @dataclass
class DownloadModelStatus(): class DownloadModelStatus():
status: str status: str
@ -38,6 +39,7 @@ class DownloadModelStatus():
"already_existed": self.already_existed "already_existed": self.already_existed
} }
async def download_model(model_download_request: Callable[[str], Awaitable[aiohttp.ClientResponse]], async def download_model(model_download_request: Callable[[str], Awaitable[aiohttp.ClientResponse]],
model_name: str, model_name: str,
model_url: str, model_url: str,
@ -114,10 +116,10 @@ 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: if os.path.commonprefix([abs_file_path, abs_base_dir]) != abs_base_dir:
raise Exception(f"Invalid model directory: {model_directory}/{model_name}") raise Exception(f"Invalid model directory: {model_directory}/{model_name}")
relative_path = '/'.join([model_directory, model_name]) relative_path = '/'.join([model_directory, model_name])
return file_path, relative_path return file_path, relative_path
async def check_file_exists(file_path: str, async def check_file_exists(file_path: str,
model_name: str, model_name: str,
progress_callback: Callable[[str, DownloadModelStatus], Awaitable[Any]], progress_callback: Callable[[str, DownloadModelStatus], Awaitable[Any]],
@ -172,6 +174,7 @@ async def track_download_progress(response: aiohttp.ClientResponse,
logging.error(traceback.format_exc()) logging.error(traceback.format_exc())
return await handle_download_error(e, model_name, progress_callback, relative_path) return await handle_download_error(e, model_name, progress_callback, relative_path)
async def handle_download_error(e: Exception, async def handle_download_error(e: Exception,
model_name: str, model_name: str,
progress_callback: Callable[[str, DownloadModelStatus], Any], progress_callback: Callable[[str, DownloadModelStatus], Any],
@ -181,6 +184,7 @@ async def handle_download_error(e: Exception,
await progress_callback(relative_path, status) await progress_callback(relative_path, status)
return status return status
def validate_model_subdirectory(model_subdirectory: str) -> bool: 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.
@ -204,6 +208,7 @@ def validate_model_subdirectory(model_subdirectory: str) -> bool:
return True return True
def validate_filename(filename: str)-> bool: def validate_filename(filename: str)-> bool:
""" """
Validate a filename to ensure it's safe and doesn't contain any path traversal attempts. Validate a filename to ensure it's safe and doesn't contain any path traversal attempts.