OPENGIT/stable-diffusion-webui
1
0
Fork 0
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui.git synced 2025-04-13 16:29:03 +08:00
Code Issues Packages Projects Releases Wiki Activity

Be more clear about Spandrel model nomenclature

Browse Source
This commit is contained in:
Aarni Koskela 2023-12-31 00:09:51 +02:00
parent c0ca6348e8
commit 777af661a2
5 changed files with 13 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
extensions-builtin/SwinIR/scripts/swinir_model.py
View File

@ -71,7 +71,7 @@ class UpscalerSwinIR(Upscaler):
else: else:
filename = path filename = path
model = modelloader.load_spandrel_model( model_descriptor = modelloader.load_spandrel_model(
filename, filename,
device=self._get_device(), device=self._get_device(),
dtype=devices.dtype, dtype=devices.dtype,
@ -79,10 +79,10 @@ class UpscalerSwinIR(Upscaler):
) )
if getattr(opts, 'SWIN_torch_compile', False): if getattr(opts, 'SWIN_torch_compile', False):
try: try:
model = torch.compile(model) model_descriptor.model.compile()
except Exception: except Exception:
logger.warning("Failed to compile SwinIR model, fallback to JIT", exc_info=True) logger.warning("Failed to compile SwinIR model, fallback to JIT", exc_info=True)
return model return model_descriptor
def _get_device(self): def _get_device(self):
return devices.get_device_for('swinir') return devices.get_device_for('swinir')

10
modules/gfpgan_model.py
View File

@ -3,6 +3,8 @@ from __future__ import annotations
import logging import logging
import os import os
import torch
from modules import ( from modules import (
devices, devices,
errors, errors,
@ -25,7 +27,7 @@ class FaceRestorerGFPGAN(face_restoration_utils.CommonFaceRestoration):
def get_device(self): def get_device(self):
return devices.device_gfpgan return devices.device_gfpgan
def load_net(self) -> None: def load_net(self) -> torch.Module:
for model_path in modelloader.load_models( for model_path in modelloader.load_models(
model_path=self.model_path, model_path=self.model_path,
model_url=model_url, model_url=model_url,
@ -34,13 +36,13 @@ 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):
net = modelloader.load_spandrel_model( model = modelloader.load_spandrel_model(
model_path, model_path,
device=self.get_device(), device=self.get_device(),
expected_architecture='GFPGAN', expected_architecture='GFPGAN',
).model ).model
net.different_w = True # see https://github.com/chaiNNer-org/spandrel/pull/81 model.different_w = True # see https://github.com/chaiNNer-org/spandrel/pull/81
return net return model
raise ValueError("No GFPGAN model found") raise ValueError("No GFPGAN model found")
def restore(self, np_image): def restore(self, np_image):

2
modules/modelloader.py
View File

@ -143,7 +143,7 @@ def load_spandrel_model(
*, *,
device: str | torch.device | None, device: str | torch.device | None,
half: bool = False, half: bool = False,
dtype: str | None = None, dtype: str | torch.dtype | None = None,
expected_architecture: str | None = None, expected_architecture: str | None = None,
) -> spandrel.ModelDescriptor: ) -> spandrel.ModelDescriptor:
import spandrel import spandrel

4
modules/realesrgan_model.py
View File

@ -36,14 +36,14 @@ class UpscalerRealESRGAN(Upscaler):
errors.report(f"Unable to load RealESRGAN model {path}", exc_info=True) errors.report(f"Unable to load RealESRGAN model {path}", exc_info=True)
return img return img
mod = modelloader.load_spandrel_model( model_descriptor = modelloader.load_spandrel_model(
info.local_data_path, info.local_data_path,
device=self.device, device=self.device,
half=(not cmd_opts.no_half and not cmd_opts.upcast_sampling), half=(not cmd_opts.no_half and not cmd_opts.upcast_sampling),
expected_architecture="ESRGAN", # "RealESRGAN" isn't a specific thing for Spandrel expected_architecture="ESRGAN", # "RealESRGAN" isn't a specific thing for Spandrel
) )
return upscale_with_model( return upscale_with_model(
mod, model_descriptor,
img, img,
tile_size=opts.ESRGAN_tile, tile_size=opts.ESRGAN_tile,
tile_overlap=opts.ESRGAN_tile_overlap, tile_overlap=opts.ESRGAN_tile_overlap,

2
modules/upscaler_utils.py
View File

@ -6,7 +6,7 @@ import torch
import tqdm import tqdm
from PIL import Image from PIL import Image
from modules import devices, images from modules import images
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)