diff --git a/README.md b/README.md index b60698a..04b228e 100644 --- a/README.md +++ b/README.md @@ -123,7 +123,14 @@ How hard ref_img guidance works: - ref_stop_step: (float) - sampling step range which enables the ref guidance mechanism (kind of scheduling) - if > 1, parse as step number; if <= 1, parse as percentage of total steps - +- upscale_* + - upscaler: (categorical) + - ratio: (float) + - width: (int) + - height: (int) + - if `width==height==0`, upscale by the specified ratio + - if `width==0 or height==0`, the zero one will be auto calculated to match the non-zero one, keeping the aspect-raio + - if `width!=0 and height!=0`, upscale while keeping the aspect-raio to cover the target size, then crop the excess if necessary ### Developers diff --git a/scripts/sonar.py b/scripts/sonar.py index 6d70ae2..595ff29 100644 --- a/scripts/sonar.py +++ b/scripts/sonar.py @@ -1,7 +1,7 @@ import os import random from PIL import Image -from typing import List +from typing import List, Tuple from pprint import pprint as pp import gradio as gr @@ -11,9 +11,12 @@ from tqdm.auto import trange from modules import scripts, devices -from modules.shared import state, opts +from modules.script_callbacks import on_before_image_saved, remove_callbacks_for_function, ImageSaveParams +from modules.shared import state, opts, sd_upscalers from modules.sd_samplers import * +from modules.ui import gr_show from modules.prompt_parser import ScheduledPromptConditioning, MulticondLearnedConditioning +from modules.images import resize_image from k_diffusion.sampling import to_d, get_ancestral_step from ldm.models.diffusion.ddpm import LatentDiffusion @@ -27,10 +30,15 @@ DEFAULT_REF_MIN_STEP = 0.0 DEFAULT_REF_MAX_STEP = 0.75 DEFAULT_REF_IMG = None +DEFAULT_UPSCALE_METH = 'Lanczos' +DEFAULT_UPSCALE_RATIO = 1.0 +DEFAULT_UPSCALE_W = 0 +DEFAULT_UPSCALE_H = 0 CHOICE_MOMENTUM_SIGN = ['pos', 'neg', 'rand'] CHOICE_MOMENTUM_HIST_INIT = ['zero', 'rand_init', 'rand_new'] CHOICE_REF_METH = ['linear', 'euler'] +CHOICE_UPSCALER = [x.name for x in sd_upscalers] # debug save latent featmap (when `Euler a`) #FEAT_MAP_PATH = 'C:\sd-webui_featmaps' @@ -813,6 +821,19 @@ def sample_img2img(self, p:StableDiffusionProcessing, x:Tensor, noise:Tensor, # ↑↑↑ the above is modified from 'modules/sd_samplers.py' ↑↑↑ +def get_upscale_resolution(p:StableDiffusionProcessing, upscale_meth:str, upscale_ratio:float, upscale_width:int, upscale_height:int) -> Tuple[bool, Tuple[int, int]]: + if upscale_meth == 'None': + return False, (p.width, p.height) + + if upscale_width == upscale_height == 0: + if upscale_ratio == 1.0: + return False, (p.width, p.height) + else: + return True, (round(p.width * upscale_ratio), round(p.height * upscale_ratio)) + else: + if upscale_width == 0: upscale_width = round(p.width * upscale_height / p.height) + if upscale_height == 0: upscale_height = round(p.height * upscale_width / p.width) + return True, (upscale_width, upscale_height) class Script(scripts.Script): @@ -829,20 +850,20 @@ def ui(self, is_img2img): sampler = gr.Radio(label='Base Sampler', value=lambda: DEFAULT_SAMPLER, choices=CHOICE_SAMPLER) with gr.Group() as tab_momentum: - with gr.Row(): + with gr.Row(variant='compact'): momentum = gr.Slider(label='Momentum (current)', minimum=0.75, maximum=1.0, value=lambda: DEFAULT_MOMENTUM) momentum_hist = gr.Slider(label='Momentum (history)', minimum=0.0, maximum=1.0, value=lambda: DEFAULT_MOMENTUM_HIST) - with gr.Row(): + with gr.Row(variant='compact'): momentum_sign = gr.Radio(label='Momentum sign', value=lambda: DEFAULT_MOMENTUM_SIGN, choices=CHOICE_MOMENTUM_SIGN) momentum_hist_init = gr.Radio(label='Momentum history init', value=lambda: DEFAULT_MOMENTUM_HIST_INIT, choices=CHOICE_MOMENTUM_HIST_INIT) with gr.Group(visible=False) as tab_file: - with gr.Row(): + with gr.Row(variant='compact'): ref_meth = gr.Radio(label='Ref guide step method', value=lambda: DEFAULT_REF_METH, choices=CHOICE_REF_METH) ref_hgf = gr.Slider(label='Ref guide factor', value=lambda: DEFAULT_REF_HGF, minimum=-1, maximum=1, step=0.001) ref_min_step = gr.Number(label='Ref start step', value=lambda: DEFAULT_REF_MIN_STEP) ref_max_step = gr.Number(label='Ref stop step', value=lambda: DEFAULT_REF_MAX_STEP) - with gr.Row(): + with gr.Row(variant='compact'): ref_img = gr.File(label='Reference image file', interactive=True) def swith_sampler(sampler:str): @@ -853,20 +874,27 @@ def swith_sampler(sampler:str): 'Naive': (True, True), } show_momt, show_file = SHOW_TABS[sampler] - return [ - { 'visible': show_momt, '__type__': 'update' }, - { 'visible': show_file, '__type__': 'update' }, + return [ + gr_show(show_momt), + gr_show(show_file), ] - sampler.change(swith_sampler, inputs=[sampler], outputs=[tab_momentum, tab_file]) + with gr.Row(variant='compact'): + upscale_meth = gr.Dropdown(label='Upscaler', value=lambda: DEFAULT_UPSCALE_METH, choices=CHOICE_UPSCALER) + upscale_ratio = gr.Slider (label='Upscale ratio', value=lambda: DEFAULT_UPSCALE_RATIO, minimum=1.0, maximum=4.0, step=0.1) + upscale_width = gr.Slider (label='Upscale width', value=lambda: DEFAULT_UPSCALE_W, minimum=0, maximum=2048, step=8) + upscale_height = gr.Slider (label='Upscale height', value=lambda: DEFAULT_UPSCALE_H, minimum=0, maximum=2048, step=8) + return [sampler, momentum, momentum_hist, momentum_hist_init, momentum_sign, - ref_meth, ref_hgf, ref_min_step, ref_max_step, ref_img] + ref_meth, ref_hgf, ref_min_step, ref_max_step, ref_img, + upscale_meth, upscale_ratio, upscale_width, upscale_height] def run(self, p:StableDiffusionProcessing, sampler:str, momentum:float, momentum_hist:float, momentum_hist_init:str, momentum_sign:str, - ref_meth:str, ref_hgf:float, ref_min_step:float, ref_max_step:float, ref_img:object): + ref_meth:str, ref_hgf:float, ref_min_step:float, ref_max_step:float, ref_img:object, + upscale_meth:str, upscale_ratio:float, upscale_width:int, upscale_height:int): # save settings to global settings['sampler'] = sampler @@ -882,7 +910,16 @@ def run(self, p:StableDiffusionProcessing, sampler:str, #pp(settings) + enabled, (tgt_w, tgt_h) = get_upscale_resolution(p, upscale_meth, upscale_ratio, upscale_width, upscale_height) + if enabled: print(f'>> upscale: ({p.width}, {p.height}) => ({tgt_w}, {tgt_h})') + + def save_image_hijack(params:ImageSaveParams): + if not enabled: return + params.image = resize_image(1, params.image, tgt_w, tgt_h, upscaler_name=upscale_meth) + state.job_count = p.n_iter * p.batch_size + on_before_image_saved(save_image_hijack) proc = process_images(p) + remove_callbacks_for_function(save_image_hijack) return proc