diff --git a/comfy_cli/command/models/models.py b/comfy_cli/command/models/models.py index 770640a0..9304e83c 100644 --- a/comfy_cli/command/models/models.py +++ b/comfy_cli/command/models/models.py @@ -1,6 +1,7 @@ import contextlib import os import pathlib +import time from typing import Annotated from urllib.parse import parse_qs, unquote, urlparse @@ -33,6 +34,18 @@ def get_workspace() -> pathlib.Path: return pathlib.Path(workspace_manager.workspace_path) +def _format_elapsed(seconds: float) -> str: + """Format elapsed seconds into a human-readable string.""" + rounded = round(seconds, 1) + if rounded < 60: + return f"{rounded:.1f}s" + minutes, secs = divmod(int(rounded), 60) + if minutes < 60: + return f"{minutes}m {secs}s" + hours, minutes = divmod(minutes, 60) + return f"{hours}h {minutes}m {secs}s" + + def potentially_strip_param_url(path_name: str) -> str: return path_name.split("?")[0] @@ -307,6 +320,8 @@ def download( print(f"[bold red]File already exists: {local_filepath}[/bold red]") return + start_time = time.monotonic() + if is_huggingface_url and check_unauthorized(url, headers): if hf_api_token is None: print( @@ -341,6 +356,9 @@ def download( print(f"Start downloading URL: {url} into {local_filepath}") download_file(url, local_filepath, headers, downloader=resolved_downloader) + elapsed = time.monotonic() - start_time + print(f"Done in {_format_elapsed(elapsed)}") + @app.command() @tracking.track_command("model") diff --git a/tests/comfy_cli/command/models/test_models.py b/tests/comfy_cli/command/models/test_models.py index dcc97df0..1f6b7134 100644 --- a/tests/comfy_cli/command/models/test_models.py +++ b/tests/comfy_cli/command/models/test_models.py @@ -4,7 +4,7 @@ import typer.testing from comfy_cli import constants -from comfy_cli.command.models.models import app, check_civitai_url, check_huggingface_url, list_models +from comfy_cli.command.models.models import _format_elapsed, app, check_civitai_url, check_huggingface_url, list_models def _make_model_tree(tmp_path: pathlib.Path) -> pathlib.Path: @@ -307,6 +307,29 @@ def test_huggingface_url_with_folder_structure(): ) +class TestFormatElapsed: + def test_under_one_minute(self): + assert _format_elapsed(5.3) == "5.3s" + + def test_fractional_seconds(self): + assert _format_elapsed(0.4) == "0.4s" + + def test_rounds_up_to_minute_boundary(self): + assert _format_elapsed(59.95) == "1m 0s" + + def test_exactly_sixty_seconds(self): + assert _format_elapsed(60) == "1m 0s" + + def test_minutes_and_seconds(self): + assert _format_elapsed(154) == "2m 34s" + + def test_over_one_hour(self): + assert _format_elapsed(3661) == "1h 1m 1s" + + def test_large_duration(self): + assert _format_elapsed(7384) == "2h 3m 4s" + + # --------------------------------------------------------------------------- # --downloader CLI option tests # --------------------------------------------------------------------------- @@ -327,7 +350,7 @@ def test_downloader_flag_forwarded(self, tmp_path): patch("comfy_cli.tracking.track_command", lambda _cmd: lambda fn: fn), ): mock_ui.prompt_input.side_effect = ["mymodel.bin", ""] - runner.invoke( + result = runner.invoke( app, [ "download", @@ -343,6 +366,7 @@ def test_downloader_flag_forwarded(self, tmp_path): assert mock_dl.called _, kwargs = mock_dl.call_args assert kwargs.get("downloader") == "aria2" + assert "Done in " in result.output def test_default_from_config(self, tmp_path): """Config default_downloader is used when no --downloader flag."""