Fix DAT models download (#16302)
Some checks failed
Linter / ruff (push) Has been cancelled
Linter / eslint (push) Has been cancelled
Tests / tests on CPU with empty model (push) Has been cancelled

This commit is contained in:
w-e-w
2024-10-24 22:05:51 +09:00
committed by GitHub
parent 5865da28d1
commit 984b952eb3
4 changed files with 97 additions and 28 deletions

View File

@@ -10,6 +10,7 @@ import torch
from modules import shared
from modules.upscaler import Upscaler, UpscalerLanczos, UpscalerNearest, UpscalerNone
from modules.util import load_file_from_url # noqa, backwards compatibility
if TYPE_CHECKING:
import spandrel
@@ -17,30 +18,6 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
def load_file_from_url(
url: str,
*,
model_dir: str,
progress: bool = True,
file_name: str | None = None,
hash_prefix: str | None = None,
) -> str:
"""Download a file from `url` into `model_dir`, using the file present if possible.
Returns the path to the downloaded file.
"""
os.makedirs(model_dir, exist_ok=True)
if not file_name:
parts = urlparse(url)
file_name = os.path.basename(parts.path)
cached_file = os.path.abspath(os.path.join(model_dir, file_name))
if not os.path.exists(cached_file):
print(f'Downloading: "{url}" to {cached_file}\n')
from torch.hub import download_url_to_file
download_url_to_file(url, cached_file, progress=progress, hash_prefix=hash_prefix)
return cached_file
def load_models(model_path: str, model_url: str = None, command_path: str = None, ext_filter=None, download_name=None, ext_blacklist=None, hash_prefix=None) -> list:
"""
A one-and done loader to try finding the desired models in specified directories.