fluxdev/stable-diffusion-webui-forge

1import cv22import torch3 4from modules_forge.shared import add_supported_preprocessor, preprocessor_dir5from ldm_patched.modules import model_management6from ldm_patched.modules.model_patcher import ModelPatcher7from modules_forge.forge_util import resize_image_with_pad8import ldm_patched.modules.clip_vision9from modules.modelloader import load_file_from_url10from modules_forge.forge_util import numpy_to_pytorch11 12 13class PreprocessorParameter:14    def __init__(self, minimum=0.0, maximum=1.0, step=0.01, label='Parameter 1', value=0.5, visible=False, **kwargs):15        self.gradio_update_kwargs = dict(16            minimum=minimum, maximum=maximum, step=step, label=label, value=value, visible=visible, **kwargs17        )18 19 20class Preprocessor:21    def __init__(self):22        self.name = 'PreprocessorBase'23        self.tags = []24        self.model_filename_filters = []25        self.slider_resolution = PreprocessorParameter(label='Resolution', minimum=128, maximum=2048, value=512, step=8, visible=True)26        self.slider_1 = PreprocessorParameter()27        self.slider_2 = PreprocessorParameter()28        self.slider_3 = PreprocessorParameter()29        self.model_patcher: ModelPatcher = None30        self.show_control_mode = True31        self.do_not_need_model = False32        self.sorting_priority = 0  # higher goes to top in the list33        self.corp_image_with_a1111_mask_when_in_img2img_inpaint_tab = True34        self.fill_mask_with_one_when_resize_and_fill = False35        self.use_soft_projection_in_hr_fix = False36        self.expand_mask_when_resize_and_fill = False37 38    def setup_model_patcher(self, model, load_device=None, offload_device=None, dtype=torch.float32, **kwargs):39        if load_device is None:40            load_device = model_management.get_torch_device()41 42        if offload_device is None:43            offload_device = torch.device('cpu')44 45        if not model_management.should_use_fp16(load_device):46            dtype = torch.float3247 48        model.eval()49        model = model.to(device=offload_device, dtype=dtype)50 51        self.model_patcher = ModelPatcher(model=model, load_device=load_device, offload_device=offload_device, **kwargs)52        self.model_patcher.dtype = dtype53        return self.model_patcher54 55    def move_all_model_patchers_to_gpu(self):56        model_management.load_models_gpu([self.model_patcher])57        return58 59    def send_tensor_to_model_device(self, x):60        return x.to(device=self.model_patcher.current_device, dtype=self.model_patcher.dtype)61 62    def process_after_running_preprocessors(self, process, params, *args, **kwargs):63        return64 65    def process_before_every_sampling(self, process, cond, mask, *args, **kwargs):66        return cond, mask67 68    def process_after_every_sampling(self, process, params, *args, **kwargs):69        return70 71    def __call__(self, input_image, resolution, slider_1=None, slider_2=None, slider_3=None, input_mask=None, **kwargs):72        return input_image73 74 75class PreprocessorNone(Preprocessor):76    def __init__(self):77        super().__init__()78        self.name = 'None'79        self.sorting_priority = 1080 81 82class PreprocessorCanny(Preprocessor):83    def __init__(self):84        super().__init__()85        self.name = 'canny'86        self.tags = ['Canny']87        self.model_filename_filters = ['canny']88        self.slider_1 = PreprocessorParameter(minimum=0, maximum=256, step=1, value=100, label='Low Threshold', visible=True)89        self.slider_2 = PreprocessorParameter(minimum=0, maximum=256, step=1, value=200, label='High Threshold', visible=True)90        self.sorting_priority = 10091        self.use_soft_projection_in_hr_fix = True92 93    def __call__(self, input_image, resolution, slider_1=None, slider_2=None, slider_3=None, **kwargs):94        input_image, remove_pad = resize_image_with_pad(input_image, resolution)95        canny_image = cv2.cvtColor(cv2.Canny(input_image, int(slider_1), int(slider_2)), cv2.COLOR_GRAY2RGB)96        return remove_pad(canny_image)97 98 99add_supported_preprocessor(PreprocessorNone())100add_supported_preprocessor(PreprocessorCanny())101 102 103class PreprocessorClipVision(Preprocessor):104    global_cache = {}105 106    def __init__(self, name, url, filename):107        super().__init__()108        self.name = name109        self.url = url110        self.filename = filename111        self.slider_resolution = PreprocessorParameter(visible=False)112        self.corp_image_with_a1111_mask_when_in_img2img_inpaint_tab = False113        self.show_control_mode = False114        self.sorting_priority = 1115        self.clipvision = None116 117    def load_clipvision(self):118        if self.clipvision is not None:119            return self.clipvision120 121        ckpt_path = load_file_from_url(122            url=self.url,123            model_dir=preprocessor_dir,124            file_name=self.filename125        )126 127        if ckpt_path in PreprocessorClipVision.global_cache:128            self.clipvision = PreprocessorClipVision.global_cache[ckpt_path]129        else:130            self.clipvision = ldm_patched.modules.clip_vision.load(ckpt_path)131            PreprocessorClipVision.global_cache[ckpt_path] = self.clipvision132 133        return self.clipvision134 135    @torch.no_grad()136    def __call__(self, input_image, resolution, slider_1=None, slider_2=None, slider_3=None, **kwargs):137        clipvision = self.load_clipvision()138        return clipvision.encode_image(numpy_to_pytorch(input_image))139