Merge pull request #16144 from akx/bump-spandrel

Bump spandrel to 0.3.4
This commit is contained in:
AUTOMATIC1111 2024-07-06 09:42:14 +03:00 committed by GitHub
commit a6c384b9f7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 32 additions and 7 deletions

View File

@ -36,13 +36,11 @@ class FaceRestorerGFPGAN(face_restoration_utils.CommonFaceRestoration):
ext_filter=['.pth'], ext_filter=['.pth'],
): ):
if 'GFPGAN' in os.path.basename(model_path): if 'GFPGAN' in os.path.basename(model_path):
model = modelloader.load_spandrel_model( return modelloader.load_spandrel_model(
model_path, model_path,
device=self.get_device(), device=self.get_device(),
expected_architecture='GFPGAN', expected_architecture='GFPGAN',
).model ).model
model.different_w = True # see https://github.com/chaiNNer-org/spandrel/pull/81
return model
raise ValueError("No GFPGAN model found") raise ValueError("No GFPGAN model found")
def restore(self, np_image): def restore(self, np_image):

View File

@ -139,6 +139,27 @@ def load_upscalers():
key=lambda x: x.name.lower() if not isinstance(x.scaler, (UpscalerNone, UpscalerLanczos, UpscalerNearest)) else "" key=lambda x: x.name.lower() if not isinstance(x.scaler, (UpscalerNone, UpscalerLanczos, UpscalerNearest)) else ""
) )
# None: not loaded, False: failed to load, True: loaded
_spandrel_extra_init_state = None
def _init_spandrel_extra_archs() -> None:
"""
Try to initialize `spandrel_extra_archs` (exactly once).
"""
global _spandrel_extra_init_state
if _spandrel_extra_init_state is not None:
return
try:
import spandrel
import spandrel_extra_arches
spandrel.MAIN_REGISTRY.add(*spandrel_extra_arches.EXTRA_REGISTRY)
_spandrel_extra_init_state = True
except Exception:
logger.warning("Failed to load spandrel_extra_arches", exc_info=True)
_spandrel_extra_init_state = False
def load_spandrel_model( def load_spandrel_model(
path: str | os.PathLike, path: str | os.PathLike,
@ -148,11 +169,16 @@ def load_spandrel_model(
dtype: str | torch.dtype | None = None, dtype: str | torch.dtype | None = None,
expected_architecture: str | None = None, expected_architecture: str | None = None,
) -> spandrel.ModelDescriptor: ) -> spandrel.ModelDescriptor:
global _spandrel_extra_init_state
import spandrel import spandrel
_init_spandrel_extra_archs()
model_descriptor = spandrel.ModelLoader(device=device).load_from_file(str(path)) model_descriptor = spandrel.ModelLoader(device=device).load_from_file(str(path))
if expected_architecture and model_descriptor.architecture != expected_architecture: arch = model_descriptor.architecture
if expected_architecture and arch.name != expected_architecture:
logger.warning( logger.warning(
f"Model {path!r} is not a {expected_architecture!r} model (got {model_descriptor.architecture!r})", f"Model {path!r} is not a {expected_architecture!r} model (got {arch.name!r})",
) )
half = False half = False
if prefer_half: if prefer_half:
@ -166,6 +192,6 @@ def load_spandrel_model(
model_descriptor.model.eval() model_descriptor.model.eval()
logger.debug( logger.debug(
"Loaded %s from %s (device=%s, half=%s, dtype=%s)", "Loaded %s from %s (device=%s, half=%s, dtype=%s)",
model_descriptor, path, device, half, dtype, arch, path, device, half, dtype,
) )
return model_descriptor return model_descriptor

View File

@ -24,7 +24,8 @@ pytorch_lightning==1.9.4
resize-right==0.0.2 resize-right==0.0.2
safetensors==0.4.2 safetensors==0.4.2
scikit-image==0.21.0 scikit-image==0.21.0
spandrel==0.1.6 spandrel==0.3.4
spandrel-extra-arches==0.1.1
tomesd==0.1.3 tomesd==0.1.3
torch torch
torchdiffeq==0.2.3 torchdiffeq==0.2.3