models.py 3.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110
  1. # Copyright 2026 The HuggingFace Team. All rights reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. """Contains commands to interact with models on the Hugging Face Hub.
  15. Usage:
  16. # list models on the Hub
  17. hf models ls
  18. # list models with a search query
  19. hf models ls --search "llama"
  20. # get info about a model
  21. hf models info Lightricks/LTX-2
  22. """
  23. import enum
  24. import json
  25. from typing import Annotated, Optional, get_args
  26. import typer
  27. from huggingface_hub.errors import RepositoryNotFoundError, RevisionNotFoundError
  28. from huggingface_hub.hf_api import ExpandModelProperty_T, ModelSort_T
  29. from huggingface_hub.utils import ANSI
  30. from ._cli_utils import (
  31. AuthorOpt,
  32. FilterOpt,
  33. LimitOpt,
  34. RevisionOpt,
  35. SearchOpt,
  36. TokenOpt,
  37. get_hf_api,
  38. make_expand_properties_parser,
  39. repo_info_to_dict,
  40. typer_factory,
  41. )
  42. _EXPAND_PROPERTIES = sorted(get_args(ExpandModelProperty_T))
  43. _SORT_OPTIONS = get_args(ModelSort_T)
  44. ModelSortEnum = enum.Enum("ModelSortEnum", {s: s for s in _SORT_OPTIONS}, type=str) # type: ignore[misc]
  45. ExpandOpt = Annotated[
  46. Optional[str],
  47. typer.Option(
  48. help=f"Comma-separated properties to expand. Example: '--expand=downloads,likes,tags'. Valid: {', '.join(_EXPAND_PROPERTIES)}.",
  49. callback=make_expand_properties_parser(_EXPAND_PROPERTIES),
  50. ),
  51. ]
  52. models_cli = typer_factory(help="Interact with models on the Hub.")
  53. @models_cli.command("ls")
  54. def models_ls(
  55. search: SearchOpt = None,
  56. author: AuthorOpt = None,
  57. filter: FilterOpt = None,
  58. sort: Annotated[
  59. Optional[ModelSortEnum],
  60. typer.Option(help="Sort results."),
  61. ] = None,
  62. limit: LimitOpt = 10,
  63. expand: ExpandOpt = None,
  64. token: TokenOpt = None,
  65. ) -> None:
  66. """List models on the Hub."""
  67. api = get_hf_api(token=token)
  68. sort_key = sort.value if sort else None
  69. results = [
  70. repo_info_to_dict(model_info)
  71. for model_info in api.list_models(
  72. filter=filter, author=author, search=search, sort=sort_key, limit=limit, expand=expand
  73. )
  74. ]
  75. print(json.dumps(results, indent=2))
  76. @models_cli.command("info")
  77. def models_info(
  78. model_id: Annotated[str, typer.Argument(help="The model ID (e.g. `username/repo-name`).")],
  79. revision: RevisionOpt = None,
  80. expand: ExpandOpt = None,
  81. token: TokenOpt = None,
  82. ) -> None:
  83. """Get info about a model on the Hub."""
  84. api = get_hf_api(token=token)
  85. try:
  86. info = api.model_info(repo_id=model_id, revision=revision, expand=expand) # type: ignore[arg-type]
  87. except RepositoryNotFoundError:
  88. print(f"Model {ANSI.bold(model_id)} not found.")
  89. raise typer.Exit(code=1)
  90. except RevisionNotFoundError:
  91. print(f"Revision {ANSI.bold(str(revision))} not found on {ANSI.bold(model_id)}.")
  92. raise typer.Exit(code=1)
  93. print(json.dumps(repo_info_to_dict(info), indent=2))