patch_hf_hub_download.py 1.5 KB

1234567891011121314151617181920212223242526272829303132333435363738394041
  1. from modules.patches import patch
  2. from modules.errors import report
  3. from inspect import signature
  4. from functools import wraps
  5. try:
  6. from huggingface_hub.utils import LocalEntryNotFoundError
  7. from huggingface_hub import file_download
  8. def try_local_files_only(func):
  9. if (param := signature(func).parameters.get('local_files_only', None)) and not param.kind == param.KEYWORD_ONLY:
  10. raise ValueError(f'{func.__name__} does not have keyword-only parameter "local_files_only"')
  11. @wraps(func)
  12. def wrapper(*args, **kwargs):
  13. try:
  14. from modules.shared import opts
  15. try_offline_mode = not kwargs.get('local_files_only') and opts.hd_dl_local_first
  16. except Exception:
  17. report('Error in try_local_files_only - skip try_local_files_only', exc_info=True)
  18. try_offline_mode = False
  19. if try_offline_mode:
  20. try:
  21. return func(*args, **{**kwargs, 'local_files_only': True})
  22. except LocalEntryNotFoundError:
  23. pass
  24. except Exception:
  25. report('Unexpected exception in try_local_files_only - retry without patch', exc_info=True)
  26. return func(*args, **kwargs)
  27. return wrapper
  28. try:
  29. patch(__name__, file_download, 'hf_hub_download', try_local_files_only(file_download.hf_hub_download))
  30. except RuntimeError:
  31. pass # already patched
  32. except Exception:
  33. report('Error patching hf_hub_download', exc_info=True)