Skip to content

Commit e522960

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 514b1b3 commit e522960

2 files changed

Lines changed: 45 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: 28 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,31 @@ def test_huggingface_url_with_folder_structure():
307307
)
308308

309309

310+
# ---------------------------------------------------------------------------
311+
# _format_elapsed tests
312+
# ---------------------------------------------------------------------------
313+
314+
315+
class TestFormatElapsed:
316+
def test_under_one_minute(self):
317+
assert _format_elapsed(5.3) == "5.3s"
318+
319+
def test_fractional_seconds(self):
320+
assert _format_elapsed(0.4) == "0.4s"
321+
322+
def test_exactly_sixty_seconds(self):
323+
assert _format_elapsed(60) == "1m 0s"
324+
325+
def test_minutes_and_seconds(self):
326+
assert _format_elapsed(154) == "2m 34s"
327+
328+
def test_over_one_hour(self):
329+
assert _format_elapsed(3661) == "1h 1m 1s"
330+
331+
def test_large_duration(self):
332+
assert _format_elapsed(7384) == "2h 3m 4s"
333+
334+
310335
# ---------------------------------------------------------------------------
311336
# --downloader CLI option tests
312337
# ---------------------------------------------------------------------------
@@ -327,7 +352,7 @@ def test_downloader_flag_forwarded(self, tmp_path):
327352
patch("comfy_cli.tracking.track_command", lambda _cmd: lambda fn: fn),
328353
):
329354
mock_ui.prompt_input.side_effect = ["mymodel.bin", ""]
330-
runner.invoke(
355+
result = runner.invoke(
331356
app,
332357
[
333358
"download",
@@ -343,6 +368,7 @@ def test_downloader_flag_forwarded(self, tmp_path):
343368
assert mock_dl.called
344369
_, kwargs = mock_dl.call_args
345370
assert kwargs.get("downloader") == "aria2"
371+
assert "Done in " in result.output
346372

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

0 commit comments

Comments
 (0)