skyadmin/cog-webui-sd

1from PIL import Image2import numpy as np3 4from modules import scripts_postprocessing, shared5import gradio as gr6 7from modules.ui_components import FormRow8 9 10upscale_cache = {}11 12 13class ScriptPostprocessingUpscale(scripts_postprocessing.ScriptPostprocessing):14    name = "Upscale"15    order = 100016 17    def ui(self):18        selected_tab = gr.State(value=0)19 20        with gr.Tabs(elem_id="extras_resize_mode"):21            with gr.TabItem('Scale by', elem_id="extras_scale_by_tab") as tab_scale_by:22                upscaling_resize = gr.Slider(minimum=1.0, maximum=8.0, step=0.05, label="Resize", value=4, elem_id="extras_upscaling_resize")23 24            with gr.TabItem('Scale to', elem_id="extras_scale_to_tab") as tab_scale_to:25                with FormRow():26                    upscaling_resize_w = gr.Number(label="Width", value=512, precision=0, elem_id="extras_upscaling_resize_w")27                    upscaling_resize_h = gr.Number(label="Height", value=512, precision=0, elem_id="extras_upscaling_resize_h")28                    upscaling_crop = gr.Checkbox(label='Crop to fit', value=True, elem_id="extras_upscaling_crop")29 30        with FormRow():31            extras_upscaler_1 = gr.Dropdown(label='Upscaler 1', elem_id="extras_upscaler_1", choices=[x.name for x in shared.sd_upscalers], value=shared.sd_upscalers[0].name)32 33        with FormRow():34            extras_upscaler_2 = gr.Dropdown(label='Upscaler 2', elem_id="extras_upscaler_2", choices=[x.name for x in shared.sd_upscalers], value=shared.sd_upscalers[0].name)35            extras_upscaler_2_visibility = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="Upscaler 2 visibility", value=0.0, elem_id="extras_upscaler_2_visibility")36 37        tab_scale_by.select(fn=lambda: 0, inputs=[], outputs=[selected_tab])38        tab_scale_to.select(fn=lambda: 1, inputs=[], outputs=[selected_tab])39 40        return {41            "upscale_mode": selected_tab,42            "upscale_by": upscaling_resize,43            "upscale_to_width": upscaling_resize_w,44            "upscale_to_height": upscaling_resize_h,45            "upscale_crop": upscaling_crop,46            "upscaler_1_name": extras_upscaler_1,47            "upscaler_2_name": extras_upscaler_2,48            "upscaler_2_visibility": extras_upscaler_2_visibility,49        }50 51    def upscale(self, image, info, upscaler, upscale_mode, upscale_by,  upscale_to_width, upscale_to_height, upscale_crop):52        if upscale_mode == 1:53            upscale_by = max(upscale_to_width/image.width, upscale_to_height/image.height)54            info["Postprocess upscale to"] = f"{upscale_to_width}x{upscale_to_height}"55        else:56            info["Postprocess upscale by"] = upscale_by57 58        cache_key = (hash(np.array(image.getdata()).tobytes()), upscaler.name, upscale_mode, upscale_by,  upscale_to_width, upscale_to_height, upscale_crop)59        cached_image = upscale_cache.pop(cache_key, None)60 61        if cached_image is not None:62            image = cached_image63        else:64            image = upscaler.scaler.upscale(image, upscale_by, upscaler.data_path)65 66        upscale_cache[cache_key] = image67        if len(upscale_cache) > shared.opts.upscaling_max_images_in_cache:68            upscale_cache.pop(next(iter(upscale_cache), None), None)69 70        if upscale_mode == 1 and upscale_crop:71            cropped = Image.new("RGB", (upscale_to_width, upscale_to_height))72            cropped.paste(image, box=(upscale_to_width // 2 - image.width // 2, upscale_to_height // 2 - image.height // 2))73            image = cropped74            info["Postprocess crop to"] = f"{image.width}x{image.height}"75 76        return image77 78    def process(self, pp: scripts_postprocessing.PostprocessedImage, upscale_mode=1, upscale_by=2.0, upscale_to_width=None, upscale_to_height=None, upscale_crop=False, upscaler_1_name=None, upscaler_2_name=None, upscaler_2_visibility=0.0):79        if upscaler_1_name == "None":80            upscaler_1_name = None81 82        upscaler1 = next(iter([x for x in shared.sd_upscalers if x.name == upscaler_1_name]), None)83        assert upscaler1 or (upscaler_1_name is None), f'could not find upscaler named {upscaler_1_name}'84 85        if not upscaler1:86            return87 88        if upscaler_2_name == "None":89            upscaler_2_name = None90 91        upscaler2 = next(iter([x for x in shared.sd_upscalers if x.name == upscaler_2_name and x.name != "None"]), None)92        assert upscaler2 or (upscaler_2_name is None), f'could not find upscaler named {upscaler_2_name}'93 94        upscaled_image = self.upscale(pp.image, pp.info, upscaler1, upscale_mode, upscale_by, upscale_to_width, upscale_to_height, upscale_crop)95        pp.info[f"Postprocess upscaler"] = upscaler1.name96 97        if upscaler2 and upscaler_2_visibility > 0:98            second_upscale = self.upscale(pp.image, pp.info, upscaler2, upscale_mode, upscale_by, upscale_to_width, upscale_to_height, upscale_crop)99            upscaled_image = Image.blend(upscaled_image, second_upscale, upscaler_2_visibility)100 101            pp.info[f"Postprocess upscaler 2"] = upscaler2.name102 103        pp.image = upscaled_image104 105    def image_changed(self):106        upscale_cache.clear()107 108 109class ScriptPostprocessingUpscaleSimple(ScriptPostprocessingUpscale):110    name = "Simple Upscale"111    order = 900112 113    def ui(self):114        with FormRow():115            upscaler_name = gr.Dropdown(label='Upscaler', choices=[x.name for x in shared.sd_upscalers], value=shared.sd_upscalers[0].name)116            upscale_by = gr.Slider(minimum=0.05, maximum=8.0, step=0.05, label="Upscale by", value=2)117 118        return {119            "upscale_by": upscale_by,120            "upscaler_name": upscaler_name,121        }122 123    def process(self, pp: scripts_postprocessing.PostprocessedImage, upscale_by=2.0, upscaler_name=None):124        if upscaler_name is None or upscaler_name == "None":125            return126 127        upscaler1 = next(iter([x for x in shared.sd_upscalers if x.name == upscaler_name]), None)128        assert upscaler1, f'could not find upscaler named {upscaler_name}'129 130        pp.image = self.upscale(pp.image, pp.info, upscaler1, 0, upscale_by, 0, 0, False)131        pp.info[f"Postprocess upscaler"] = upscaler1.name132