diff --git a/modules/dat_model.py b/modules/dat_model.py index cccc017ff..298d160d1 100644 --- a/modules/dat_model.py +++ b/modules/dat_model.py @@ -51,6 +51,16 @@ class UpscalerDAT(Upscaler): model_dir=self.model_download_path, hash_prefix=scaler.sha256, ) + + if os.path.getsize(scaler.local_data_path) < 200: + # Re-download if the file is too small, probably an LFS pointer + scaler.local_data_path = modelloader.load_file_from_url( + scaler.data_path, + model_dir=self.model_download_path, + hash_prefix=scaler.sha256, + re_download=True, + ) + if not os.path.exists(scaler.local_data_path): raise FileNotFoundError(f"DAT data missing: {scaler.local_data_path}") return scaler diff --git a/modules/modelloader.py b/modules/modelloader.py index 36e7415af..1596e11b0 100644 --- a/modules/modelloader.py +++ b/modules/modelloader.py @@ -24,17 +24,22 @@ def load_file_from_url( progress: bool = True, file_name: str | None = None, hash_prefix: str | None = None, + re_download: bool = False, ) -> str: """Download a file from `url` into `model_dir`, using the file present if possible. - Returns the path to the downloaded file. + + file_name: if specified, it will be used as the filename, otherwise the filename will be extracted from the url. + hash_prefix: is provided, the hash of the downloaded file will be checked against. + re_download: re-download the file even if it already exists. + """ 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): + if re_download or 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)