modelloader.py 6.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169
  1. from __future__ import annotations
  2. import importlib
  3. import logging
  4. import os
  5. from typing import TYPE_CHECKING
  6. from urllib.parse import urlparse
  7. import torch
  8. from modules import shared
  9. from modules.upscaler import Upscaler, UpscalerLanczos, UpscalerNearest, UpscalerNone
  10. if TYPE_CHECKING:
  11. import spandrel
  12. logger = logging.getLogger(__name__)
  13. def load_file_from_url(
  14. url: str,
  15. *,
  16. model_dir: str,
  17. progress: bool = True,
  18. file_name: str | None = None,
  19. ) -> str:
  20. """Download a file from `url` into `model_dir`, using the file present if possible.
  21. Returns the path to the downloaded file.
  22. """
  23. os.makedirs(model_dir, exist_ok=True)
  24. if not file_name:
  25. parts = urlparse(url)
  26. file_name = os.path.basename(parts.path)
  27. cached_file = os.path.abspath(os.path.join(model_dir, file_name))
  28. if not os.path.exists(cached_file):
  29. print(f'Downloading: "{url}" to {cached_file}\n')
  30. from torch.hub import download_url_to_file
  31. download_url_to_file(url, cached_file, progress=progress)
  32. return cached_file
  33. def load_models(model_path: str, model_url: str = None, command_path: str = None, ext_filter=None, download_name=None, ext_blacklist=None) -> list:
  34. """
  35. A one-and done loader to try finding the desired models in specified directories.
  36. @param download_name: Specify to download from model_url immediately.
  37. @param model_url: If no other models are found, this will be downloaded on upscale.
  38. @param model_path: The location to store/find models in.
  39. @param command_path: A command-line argument to search for models in first.
  40. @param ext_filter: An optional list of filename extensions to filter by
  41. @return: A list of paths containing the desired model(s)
  42. """
  43. output = []
  44. try:
  45. places = []
  46. if command_path is not None and command_path != model_path:
  47. pretrained_path = os.path.join(command_path, 'experiments/pretrained_models')
  48. if os.path.exists(pretrained_path):
  49. print(f"Appending path: {pretrained_path}")
  50. places.append(pretrained_path)
  51. elif os.path.exists(command_path):
  52. places.append(command_path)
  53. places.append(model_path)
  54. for place in places:
  55. for full_path in shared.walk_files(place, allowed_extensions=ext_filter):
  56. if os.path.islink(full_path) and not os.path.exists(full_path):
  57. print(f"Skipping broken symlink: {full_path}")
  58. continue
  59. if ext_blacklist is not None and any(full_path.endswith(x) for x in ext_blacklist):
  60. continue
  61. if full_path not in output:
  62. output.append(full_path)
  63. if model_url is not None and len(output) == 0:
  64. if download_name is not None:
  65. output.append(load_file_from_url(model_url, model_dir=places[0], file_name=download_name))
  66. else:
  67. output.append(model_url)
  68. except Exception:
  69. pass
  70. return output
  71. def friendly_name(file: str):
  72. if file.startswith("http"):
  73. file = urlparse(file).path
  74. file = os.path.basename(file)
  75. model_name, extension = os.path.splitext(file)
  76. return model_name
  77. def load_upscalers():
  78. # We can only do this 'magic' method to dynamically load upscalers if they are referenced,
  79. # so we'll try to import any _model.py files before looking in __subclasses__
  80. modules_dir = os.path.join(shared.script_path, "modules")
  81. for file in os.listdir(modules_dir):
  82. if "_model.py" in file:
  83. model_name = file.replace("_model.py", "")
  84. full_model = f"modules.{model_name}_model"
  85. try:
  86. importlib.import_module(full_model)
  87. except Exception:
  88. pass
  89. data = []
  90. commandline_options = vars(shared.cmd_opts)
  91. # some of upscaler classes will not go away after reloading their modules, and we'll end
  92. # up with two copies of those classes. The newest copy will always be the last in the list,
  93. # so we go from end to beginning and ignore duplicates
  94. used_classes = {}
  95. for cls in reversed(Upscaler.__subclasses__()):
  96. classname = str(cls)
  97. if classname not in used_classes:
  98. used_classes[classname] = cls
  99. for cls in reversed(used_classes.values()):
  100. name = cls.__name__
  101. cmd_name = f"{name.lower().replace('upscaler', '')}_models_path"
  102. commandline_model_path = commandline_options.get(cmd_name, None)
  103. scaler = cls(commandline_model_path)
  104. scaler.user_path = commandline_model_path
  105. scaler.model_download_path = commandline_model_path or scaler.model_path
  106. data += scaler.scalers
  107. shared.sd_upscalers = sorted(
  108. data,
  109. # Special case for UpscalerNone keeps it at the beginning of the list.
  110. key=lambda x: x.name.lower() if not isinstance(x.scaler, (UpscalerNone, UpscalerLanczos, UpscalerNearest)) else ""
  111. )
  112. def load_spandrel_model(
  113. path: str | os.PathLike,
  114. *,
  115. device: str | torch.device | None,
  116. prefer_half: bool = False,
  117. dtype: str | torch.dtype | None = None,
  118. expected_architecture: str | None = None,
  119. ) -> spandrel.ModelDescriptor:
  120. import spandrel
  121. model_descriptor = spandrel.ModelLoader(device=device).load_from_file(str(path))
  122. if expected_architecture and model_descriptor.architecture != expected_architecture:
  123. logger.warning(
  124. f"Model {path!r} is not a {expected_architecture!r} model (got {model_descriptor.architecture!r})",
  125. )
  126. half = False
  127. if prefer_half:
  128. if model_descriptor.supports_half:
  129. model_descriptor.model.half()
  130. half = True
  131. else:
  132. logger.info("Model %s does not support half precision, ignoring --half", path)
  133. if dtype:
  134. model_descriptor.model.to(dtype=dtype)
  135. model_descriptor.model.eval()
  136. logger.debug(
  137. "Loaded %s from %s (device=%s, half=%s, dtype=%s)",
  138. model_descriptor, path, device, half, dtype,
  139. )
  140. return model_descriptor