Commit b6e5edd7 authored by AUTOMATIC's avatar AUTOMATIC
Browse files

add built-in extension system

add support for adding upscalers in extensions
move LDSR, ScuNET and SwinIR to built-in extensions
parent 46b0d230
Loading
Loading
Loading
Loading
+6 −0
Original line number Diff line number Diff line
import os
from modules import paths


def preload(parser):
    parser.add_argument("--ldsr-models-path", type=str, help="Path to directory with LDSR model file(s).", default=os.path.join(paths.models_path, 'LDSR'))
+11 −2
Original line number Diff line number Diff line
@@ -5,8 +5,8 @@ import traceback
from basicsr.utils.download_util import load_file_from_url

from modules.upscaler import Upscaler, UpscalerData
from modules.ldsr_model_arch import LDSR
from modules import shared
from ldsr_model_arch import LDSR
from modules import shared, script_callbacks


class UpscalerLDSR(Upscaler):
@@ -52,3 +52,12 @@ class UpscalerLDSR(Upscaler):
            return img
        ddim_steps = shared.opts.ldsr_steps
        return ldsr.super_resolution(img, ddim_steps, self.scale)


def on_ui_settings():
    import gradio as gr

    shared.opts.add_option("ldsr_steps", shared.OptionInfo(100, "LDSR processing steps. Lower = faster", gr.Slider, {"minimum": 1, "maximum": 200, "step": 1}, section=('upscaling', "Upscaling")))


script_callbacks.on_ui_settings(on_ui_settings)
+6 −0
Original line number Diff line number Diff line
import os
from modules import paths


def preload(parser):
    parser.add_argument("--scunet-models-path", type=str, help="Path to directory with ScuNET model file(s).", default=os.path.join(paths.models_path, 'ScuNET'))
+3 −3
Original line number Diff line number Diff line
@@ -9,7 +9,7 @@ from basicsr.utils.download_util import load_file_from_url

import modules.upscaler
from modules import devices, modelloader
from modules.scunet_model_arch import SCUNet as net
from scunet_model_arch import SCUNet as net


class UpscalerScuNET(modules.upscaler.Upscaler):
@@ -49,7 +49,7 @@ class UpscalerScuNET(modules.upscaler.Upscaler):
        if model is None:
            return img

        device = devices.device_scunet
        device = devices.get_device_for('scunet')
        img = np.array(img)
        img = img[:, :, ::-1]
        img = np.moveaxis(img, 2, 0) / 255
@@ -66,7 +66,7 @@ class UpscalerScuNET(modules.upscaler.Upscaler):
        return PIL.Image.fromarray(output, 'RGB')

    def load_model(self, path: str):
        device = devices.device_scunet
        device = devices.get_device_for('scunet')
        if "http" in path:
            filename = load_file_from_url(url=self.model_url, model_dir=self.model_path, file_name="%s.pth" % self.name,
                                          progress=True)
Loading