Skip to content

Commit 76865c9

Browse files
committed
feat: show elapsed time after model downloads
Print "Done in Xs" / "Xm Ys" / "Xh Ym Zs" after successful model downloads, covering both the httpx/aria2 and HuggingFace Hub paths. Closes #421
1 parent b414616 commit 76865c9

2 files changed

Lines changed: 41 additions & 2 deletions

File tree

comfy_cli/command/models/models.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import contextlib
22
import os
33
import pathlib
4+
import time
45
from typing import Annotated
56
from urllib.parse import parse_qs, unquote, urlparse
67

@@ -33,6 +34,17 @@ def get_workspace() -> pathlib.Path:
3334
return pathlib.Path(workspace_manager.workspace_path)
3435

3536

37+
def _format_elapsed(seconds: float) -> str:
38+
"""Format elapsed seconds into a human-readable string."""
39+
if seconds < 60:
40+
return f"{seconds:.1f}s"
41+
minutes, secs = divmod(int(seconds), 60)
42+
if minutes < 60:
43+
return f"{minutes}m {secs}s"
44+
hours, minutes = divmod(minutes, 60)
45+
return f"{hours}h {minutes}m {secs}s"
46+
47+
3648
def potentially_strip_param_url(path_name: str) -> str:
3749
return path_name.split("?")[0]
3850

@@ -307,6 +319,8 @@ def download(
307319
print(f"[bold red]File already exists: {local_filepath}[/bold red]")
308320
return
309321

322+
start_time = time.monotonic()
323+
310324
if is_huggingface_url and check_unauthorized(url, headers):
311325
if hf_api_token is None:
312326
print(
@@ -341,6 +355,9 @@ def download(
341355
print(f"Start downloading URL: {url} into {local_filepath}")
342356
download_file(url, local_filepath, headers, downloader=resolved_downloader)
343357

358+
elapsed = time.monotonic() - start_time
359+
print(f"Done in {_format_elapsed(elapsed)}")
360+
344361

345362
@app.command()
346363
@tracking.track_command("model")

tests/comfy_cli/command/models/test_models.py

Lines changed: 24 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
import typer.testing
55

66
from comfy_cli import constants
7-
from comfy_cli.command.models.models import app, check_civitai_url, check_huggingface_url, list_models
7+
from comfy_cli.command.models.models import _format_elapsed, app, check_civitai_url, check_huggingface_url, list_models
88

99

1010
def _make_model_tree(tmp_path: pathlib.Path) -> pathlib.Path:
@@ -307,6 +307,27 @@ def test_huggingface_url_with_folder_structure():
307307
)
308308

309309

310+
311+
class TestFormatElapsed:
312+
def test_under_one_minute(self):
313+
assert _format_elapsed(5.3) == "5.3s"
314+
315+
def test_fractional_seconds(self):
316+
assert _format_elapsed(0.4) == "0.4s"
317+
318+
def test_exactly_sixty_seconds(self):
319+
assert _format_elapsed(60) == "1m 0s"
320+
321+
def test_minutes_and_seconds(self):
322+
assert _format_elapsed(154) == "2m 34s"
323+
324+
def test_over_one_hour(self):
325+
assert _format_elapsed(3661) == "1h 1m 1s"
326+
327+
def test_large_duration(self):
328+
assert _format_elapsed(7384) == "2h 3m 4s"
329+
330+
310331
# ---------------------------------------------------------------------------
311332
# --downloader CLI option tests
312333
# ---------------------------------------------------------------------------
@@ -327,7 +348,7 @@ def test_downloader_flag_forwarded(self, tmp_path):
327348
patch("comfy_cli.tracking.track_command", lambda _cmd: lambda fn: fn),
328349
):
329350
mock_ui.prompt_input.side_effect = ["mymodel.bin", ""]
330-
runner.invoke(
351+
result = runner.invoke(
331352
app,
332353
[
333354
"download",
@@ -343,6 +364,7 @@ def test_downloader_flag_forwarded(self, tmp_path):
343364
assert mock_dl.called
344365
_, kwargs = mock_dl.call_args
345366
assert kwargs.get("downloader") == "aria2"
367+
assert "Done in " in result.output
346368

347369
def test_default_from_config(self, tmp_path):
348370
"""Config default_downloader is used when no --downloader flag."""

0 commit comments

Comments
 (0)