[Various Changes] GoBig fixes, model loading unloading and more (#553)

* added image lab

* first release

model loading/unloading and save procedure added, commented out unused code from frontend

* bug fixes

Changed the image output to a gallery to display multiple items

Fixed results not showing up in output

Fixed RealESRGAN 2x mode not working and hard coded the default value for the reload.

* added GoBig model check

* added LDSR load check

* removed global statements, added model loader/unloader function

* fixed optimized mode

* update

* update

Added send to lab button
Added a print out if latent-diffusion folder isn't found

* brought back the fix faces and upscale in generation tab

* uncommenting img lab flag

* added LDSR instructions

* default imgProcessorTask set to false

* exposed LDSR settings to lab

users need to reclone the LDSR repo to use them.

* Update frontend.py

moving some stuff around to make them more coherent

* restored upscale and fix faces to img2img

* added notice section

* fixed gfpgan/upscaled pictures not showing in 2img interfaces

* send to lab button now sends info as well

* uncommented dimension info update

* added increment buttons to sampler for that k_euler_a action

* image lab settings toggle on and off with selection

* removed wip settings panel

* better model loading handling and removed increment buttons

* explaining

* disabled SD unloading in image lab upscaling with realesgan and face fix

* fixed a conflict with image lab

Co-authored-by: dr3amer <91037083+dr3am37@users.noreply.github.com>
Co-authored-by: hlky <106811348+hlky@users.noreply.github.com>
This commit is contained in:
devilismyfriend 2022-09-03 01:07:17 -07:00 committed by GitHub
parent b61b453a02
commit de95a09593
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 259 additions and 102 deletions

View File

@ -2,6 +2,9 @@
# This repo is for development, there may be bugs and new features
# Notice
-New LDSR settings added to Image Lab, To use the new LDSR settings please make sure to re-clone the LDSR (Instructions added below) to insure you have the latest.
## Feature request? Use [discussions](https://github.com/hlky/stable-diffusion-webui/discussions)
### Questions about **_[Upscalers](https://github.com/hlky/stable-diffusion-webui/wiki/Upscalers)_**?
@ -62,6 +65,15 @@ into the `/stable-diffusion/src/gfpgan/experiments/pretrained_models` directory.
Download [RealESRGAN_x4plus.pth](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth) and [RealESRGAN_x4plus_anime_6B.pth](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth).
Put them into the `stable-diffusion/src/realesrgan/experiments/pretrained_models` directory.
### LDSR
Quadruple your resolution using Latent Diffusion, to install:
- Git clone https://github.com/devilismyfriend/latent-diffusion into your stable-diffusion-main/src/ folder
- Rename latent-diffusion-main folder to latent-diffusion
- If on windows: run download_models.bat to download the required model files
- Otherwise to manually install the model download [project.yaml](https://heibox.uni-heidelberg.de/f/31a76b13ea27482981b4/?dl=1) and [last.cpkt](https://heibox.uni-heidelberg.de/f/578df07c8fc04ffbadf3/?dl=1) and rename last.ckpt to model.ckpt
- Place both under stable-diffusion-main/src/latent-diffusion/experiments/pretrained_models/
- Make sure you have both project.yaml and model.ckpt in that folder and path.
- LDSR should be wokring now.
### Web UI
When launching, you may get a very long warning message related to some weights not being used. You may freely ignore it.
@ -87,6 +99,7 @@ also a separate tab that just allows you to use GFPGAN on any picture, with a sl
Lets you double the resolution of generated images. There is a checkbox in every tab to use RealESRGAN, and you can choose between the regular upscaler and the anime version.
There is also a separate tab for using RealESRGAN on any picture.
![](images/RealESRGAN.png)
### Sampling method selection

View File

@ -4,7 +4,16 @@
max-width:89vw
}
#increment_btn_minus {
align-self: center;
background: none;
border: none;
}
#increment_btn_plus {
align-self: center;
background: none;
border: none;
}
#prompt_input, #img2img_prompt_input {
padding: 0px;
border: none;

View File

@ -1,3 +1,4 @@
from email.policy import default
import sys
from tkinter.filedialog import askopenfilename
import gradio as gr
@ -77,9 +78,15 @@ def draw_gradio_ui(opt, img2img=lambda x: x, txt2img=lambda x: x,imgproc=lambda
input_txt2img_defaults = gr.Button('Restore defaults')
output_txt2img_stats = gr.HTML(label='Stats')
with gr.Column():
with gr.Row():
#Commenting out incrementing/decrementing buttons for now
#increment_btn_minus = gr.Button("-", elem_id="increment_btn_minus")
txt2img_steps = gr.Slider(minimum=1, maximum=250, step=1, label="Sampling Steps",
value=txt2img_defaults['ddim_steps'])
#increment_btn_plus = gr.Button("+", elem_id="increment_btn_plus")
#increment_btn_minus.click(fn=uifn.increment_down,inputs=[txt2img_steps], outputs=[txt2img_steps])
#increment_btn_plus.click(fn=uifn.increment_up,inputs=[txt2img_steps], outputs=[txt2img_steps])
txt2img_steps = gr.Slider(minimum=1, maximum=250, step=1, label="Sampling Steps",
value=txt2img_defaults['ddim_steps'])
txt2img_sampling = gr.Dropdown(label='Sampling method (k_lms is default k-diffusion sampler)',
choices=["DDIM", "PLMS", 'k_dpm_2_a', 'k_dpm_2', 'k_euler_a',
'k_euler', 'k_heun', 'k_lms'],
@ -139,6 +146,10 @@ def draw_gradio_ui(opt, img2img=lambda x: x, txt2img=lambda x: x,imgproc=lambda
txt2img_inputs,
txt2img_outputs
)
txt2img_width.change(fn=uifn.update_dimensions_info, inputs=[txt2img_width, txt2img_height], outputs=txt2img_dimensions_info_text_box)
txt2img_height.change(fn=uifn.update_dimensions_info, inputs=[txt2img_width, txt2img_height], outputs=txt2img_dimensions_info_text_box)
txt2img_settings_elements = [
txt2img_prompt, txt2img_steps, txt2img_sampling, txt2img_toggles, txt2img_realesrgan_model_name,
txt2img_ddim_eta, txt2img_batch_count, txt2img_batch_size, txt2img_cfg, txt2img_seed,
@ -170,6 +181,7 @@ def draw_gradio_ui(opt, img2img=lambda x: x, txt2img=lambda x: x,imgproc=lambda
# txt2img_width.change(fn=uifn.update_dimensions_info, inputs=[txt2img_width, txt2img_height], outputs=txt2img_dimensions_info_text_box)
# txt2img_height.change(fn=uifn.update_dimensions_info, inputs=[txt2img_width, txt2img_height], outputs=txt2img_dimensions_info_text_box)
live_prompt_params = [txt2img_prompt, txt2img_width, txt2img_height, txt2img_steps, txt2img_seed, txt2img_batch_count, txt2img_cfg]
txt2img_prompt.change(
fn=None,
@ -412,47 +424,72 @@ def draw_gradio_ui(opt, img2img=lambda x: x, txt2img=lambda x: x,imgproc=lambda
with gr.TabItem('Output'):
imgproc_output = gr.Gallery(label="Output", elem_id="imgproc_gallery_output")
with gr.Row(elem_id="proc_options_row"):
imgproc_toggles = gr.CheckboxGroup(label='Processor Modes', choices=imgproc_mode_toggles, type="index")
with gr.Tabs():
with gr.TabItem('Fix Face Settings'):
gfpgan_defaults = {
'strength': 100,
}
with gr.Box():
with gr.Column():
gr.Markdown("<b>Processor Selection</b>")
imgproc_toggles = gr.CheckboxGroup(label = '',choices=imgproc_mode_toggles, type="index")
#.change toggles to show options
#imgproc_toggles.change()
with gr.Box(visible=False) as gfpgan_group:
if 'gfpgan' in user_defaults:
gfpgan_defaults.update(user_defaults['gfpgan'])
if GFPGAN is None:
gr.HTML("""
<div id="90" style="max-width: 100%; font-size: 14px; text-align: center;" class="output-markdown gr-prose border-solid border border-gray-200 rounded gr-panel">
<p><b> Please download GFPGAN to activate face fixing features</b>, instructions are available at the <a href='https://github.com/hlky/stable-diffusion-webui'>Github</a></p>
</div>
""")
#gr.Markdown("")
#gr.Markdown("<b> Please download GFPGAN to activate face fixing features</b>, instructions are available at the <a href='https://github.com/hlky/stable-diffusion-webui'>Github</a>")
gfpgan_defaults = {
'strength': 100,
}
if 'gfpgan' in user_defaults:
gfpgan_defaults.update(user_defaults['gfpgan'])
if GFPGAN is None:
gr.HTML("""
<div id="90" style="max-width: 100%; font-size: 14px; text-align: center;" class="output-markdown gr-prose border-solid border border-gray-200 rounded gr-panel">
<p><b> Please download GFPGAN to activate face fixing features</b>, instructions are available at the <a href='https://github.com/hlky/stable-diffusion-webui'>Github</a></p>
</div>
""")
#gr.Markdown("")
#gr.Markdown("<b> Please download GFPGAN to activate face fixing features</b>, instructions are available at the <a href='https://github.com/hlky/stable-diffusion-webui'>Github</a>")
with gr.Column():
gr.Markdown("<b>GFPGAN Settings</b>")
imgproc_gfpgan_strength = gr.Slider(minimum=0.0, maximum=1.0, step=0.001, label="Effect strength",
value=gfpgan_defaults['strength'],visible=GFPGAN is not None)
with gr.TabItem('Upscale Settings'):
imgproc_realesrgan_model_name = gr.Dropdown(label='RealESRGAN model', interactive=RealESRGAN is not None,
choices= ['RealESRGAN_x4plus',
'RealESRGAN_x4plus_anime_6B','RealESRGAN_x2plus',
'RealESRGAN_x2plus_anime_6B'],
value='RealESRGAN_x4plus',
visible=RealESRGAN is not None) # TODO: Feels like I shouldnt slot it in here.
if LDSR:
upscaleModes = ['RealESRGAN','GoBig','Latent Diffusion SR','GoLatent ']
else:
gr.HTML("""
<div id="90" style="max-width: 100%; font-size: 14px; text-align: center;" class="output-markdown gr-prose border-solid border border-gray-200 rounded gr-panel">
<p><b> Please download LDSR to activate more upscale features</b>, instructions are available at the <a href='https://github.com/hlky/stable-diffusion-webui'>Github</a></p>
</div>
""")
upscaleModes = ['RealESRGAN','GoBig','Latent Diffusion SR']
#gr.Markdown("<b> Please download LDSR to activate more upscale features</b>, instructions are available at the <a href='https://github.com/hlky/stable-diffusion-webui'>Github</a>")
upscaleModes = ['RealESRGAN','GoBig']
imgproc_upscale_toggles = gr.Radio(label='Upscale Modes', choices=upscaleModes, type="index",visible=RealESRGAN is not None)
with gr.Box(visible=False) as upscale_group:
if LDSR:
upscaleModes = ['RealESRGAN','GoBig','Latent Diffusion SR','GoLatent ']
else:
gr.HTML("""
<div id="90" style="max-width: 100%; font-size: 14px; text-align: center;" class="output-markdown gr-prose border-solid border border-gray-200 rounded gr-panel">
<p><b> Please download LDSR to activate more upscale features</b>, instructions are available at the <a href='https://github.com/hlky/stable-diffusion-webui'>Github</a></p>
</div>
""")
upscaleModes = ['RealESRGAN','GoBig']
with gr.Column():
gr.Markdown("<b>Upscaler Selection</b>")
imgproc_upscale_toggles = gr.Radio(label = '',choices=upscaleModes, type="index",visible=RealESRGAN is not None,value='RealESRGAN')
with gr.Box(visible=False) as upscalerSettings_group:
with gr.Box(visible=True) as realesrgan_group:
with gr.Column():
gr.Markdown("<b>RealESRGAN Settings</b>")
imgproc_realesrgan_model_name = gr.Dropdown(label='RealESRGAN model', interactive=RealESRGAN is not None,
choices= ['RealESRGAN_x4plus',
'RealESRGAN_x4plus_anime_6B','RealESRGAN_x2plus',
'RealESRGAN_x2plus_anime_6B'],
value='RealESRGAN_x4plus',
visible=RealESRGAN is not None) # TODO: Feels like I shouldnt slot it in here.
with gr.Box(visible=False) as ldsr_group:
with gr.Row(elem_id="ldsr_settings_row"):
with gr.Column():
gr.Markdown("<b>Latent Diffusion Super Sampling Settings</b>")
imgproc_ldsr_steps = gr.Slider(minimum=0, maximum=500, step=10, label="LDSR Sampling Steps",
value=100,visible=LDSR is not None)
imgproc_ldsr_pre_downSample = gr.Dropdown(label='LDSR Pre Downsample mode (Lower resolution before processing for speed)',
choices=["None", '1/2', '1/4'],value="None",visible=LDSR is not None)
imgproc_ldsr_post_downSample = gr.Dropdown(label='LDSR Post Downsample mode (aka SuperSampling)',
choices=["None", "Original Size", '1/2', '1/4'],value="None",visible=LDSR is not None)
with gr.Box(visible=False) as gobig_group:
with gr.Row(elem_id="proc_prompt_row"):
with gr.Column():
imgproc_prompt = gr.Textbox(label="These settings are applied only for GoBig and GoLatent modes",
gr.Markdown("<b>GoBig Settings</b>")
imgproc_prompt = gr.Textbox(label="",
elem_id='prompt_input',
placeholder="A corgi wearing a top hat as an oil painting.",
lines=1,
@ -480,7 +517,8 @@ def draw_gradio_ui(opt, img2img=lambda x: x, txt2img=lambda x: x,imgproc=lambda
imgproc_btn.click(
imgproc,
[imgproc_source, imgproc_folder,imgproc_prompt,imgproc_toggles,
imgproc_upscale_toggles,imgproc_realesrgan_model_name,imgproc_sampling, imgproc_steps, imgproc_height, imgproc_width, imgproc_cfg, imgproc_denoising, imgproc_seed,imgproc_gfpgan_strength],
imgproc_upscale_toggles,imgproc_realesrgan_model_name,imgproc_sampling, imgproc_steps, imgproc_height,
imgproc_width, imgproc_cfg, imgproc_denoising, imgproc_seed,imgproc_gfpgan_strength,imgproc_ldsr_steps,imgproc_ldsr_pre_downSample,imgproc_ldsr_post_downSample],
[imgproc_output])
imgproc_source.change(
@ -489,9 +527,15 @@ def draw_gradio_ui(opt, img2img=lambda x: x, txt2img=lambda x: x,imgproc=lambda
[imgproc_pngnfo] )
output_txt2img_to_imglab.click(
uifn.copy_img_to_lab,
[output_txt2img_gallery],
[imgproc_source, tabs],
fn=uifn.copy_img_params_to_lab,
inputs = [output_txt2img_params],
outputs = [imgproc_prompt,imgproc_seed,imgproc_steps,imgproc_cfg,imgproc_sampling],
)
output_txt2img_to_imglab.click(
fn=uifn.copy_img_to_lab,
inputs = [output_txt2img_gallery],
outputs = [imgproc_source, tabs],
_js=call_JS("moveImageFromGallery",
fromId="txt2img_gallery_output",
toId="imglab_input")
@ -505,7 +549,13 @@ def draw_gradio_ui(opt, img2img=lambda x: x, txt2img=lambda x: x,imgproc=lambda
<p><b> Please download RealESRGAN to activate upscale features</b>, instructions are available at the <a href='https://github.com/hlky/stable-diffusion-webui'>Github</a></p>
</div>
""")
imgproc_toggles.change(fn=uifn.toggle_options_gfpgan, inputs=[imgproc_toggles], outputs=[gfpgan_group])
imgproc_toggles.change(fn=uifn.toggle_options_upscalers, inputs=[imgproc_toggles], outputs=[upscale_group])
imgproc_toggles.change(fn=uifn.toggle_options_upscalers, inputs=[imgproc_toggles], outputs=[upscalerSettings_group])
imgproc_upscale_toggles.change(fn=uifn.toggle_options_realesrgan, inputs=[imgproc_upscale_toggles], outputs=[realesrgan_group])
imgproc_upscale_toggles.change(fn=uifn.toggle_options_ldsr, inputs=[imgproc_upscale_toggles], outputs=[ldsr_group])
imgproc_upscale_toggles.change(fn=uifn.toggle_options_gobig, inputs=[imgproc_upscale_toggles], outputs=[gobig_group])
"""
if GFPGAN is not None:
gfpgan_defaults = {

View File

@ -22,16 +22,64 @@ def update_image_mask(cropped_image, resize_mode, width, height):
resized_cropped_image = resize_image(resize_mode, cropped_image, width, height) if cropped_image else None
return gr.update(value=resized_cropped_image)
def toggle_options_gfpgan(selection):
if 0 in selection:
return gr.update(visible=True)
else:
return gr.update(visible=False)
def toggle_options_upscalers(selection):
if 1 in selection:
return gr.update(visible=True)
else:
return gr.update(visible=False)
def toggle_options_realesrgan(selection):
if selection == 0 or selection == 1 or selection == 3:
return gr.update(visible=True)
else:
return gr.update(visible=False)
def toggle_options_gobig(selection):
if selection == 1:
#print(selection)
return gr.update(visible=True)
if selection == 3:
return gr.update(visible=True)
else:
return gr.update(visible=False)
def toggle_options_ldsr(selection):
if selection == 2 or selection == 3:
return gr.update(visible=True)
else:
return gr.update(visible=False)
def increment_down(value):
return value - 1
def increment_up(value):
return value + 1
def copy_img_to_lab(img):
try:
image_data = re.sub('^data:image/.+;base64,', '', img)
processed_image = Image.open(BytesIO(base64.b64decode(image_data)))
tab_update = gr.update(selected='imgproc_tab')
img_update = gr.update(value=processed_image)
return processed_image, tab_update
return processed_image, tab_update,
except IndexError:
return [None, None]
def copy_img_params_to_lab(params):
try:
prompt = params[0][0].replace('\n', ' ').replace('\r', '')
seed = int(params[1][1])
steps = int(params[7][1])
cfg_scale = float(params[9][1])
sampler = params[11][1]
return prompt,seed,steps,cfg_scale,sampler
except IndexError:
return [None, None]
def copy_img_to_input(img):
try:
image_data = re.sub('^data:image/.+;base64,', '', img)
@ -128,7 +176,6 @@ def update_dimensions_info(width, height):
pixel_count_formated = "{:,.0f}".format(width * height)
return f"Aspect ratio: {round(width / height, 5)}\nTotal pixel count: {pixel_count_formated}"
def get_png_nfo( image: Image ):
info_text = ""
visible = bool(image and any(image.info))

148
webui.py
View File

@ -171,9 +171,15 @@ def crash(e, s):
global device
print(s, '\n', e)
del model
del device
try:
del model
del device
except:
try:
del device
except:
pass
pass
print('exiting...calling os._exit(0)')
t = threading.Timer(0.25, os._exit, args=[0])
@ -282,7 +288,7 @@ def create_random_tensors(shape, seeds):
def torch_gc():
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
def load_LDSR():
def load_LDSR(checking=False):
model_name = 'model'
yaml_name = 'project'
model_path = os.path.join(LDSR_dir, 'experiments/pretrained_models', model_name + '.ckpt')
@ -291,17 +297,20 @@ def load_LDSR():
raise Exception("LDSR model not found at path "+model_path)
if not os.path.isfile(yaml_path):
raise Exception("LDSR model not found at path "+yaml_path)
if checking == True:
return True
sys.path.append(os.path.abspath(LDSR_dir))
from LDSR import LDSR
LDSRObject = LDSR(model_path, yaml_path)
return LDSRObject
def load_GFPGAN():
def load_GFPGAN(checking=False):
model_name = 'GFPGANv1.3'
model_path = os.path.join(GFPGAN_dir, 'experiments/pretrained_models', model_name + '.pth')
if not os.path.isfile(model_path):
raise Exception("GFPGAN model not found at path "+model_path)
if checking == True:
return True
sys.path.append(os.path.abspath(GFPGAN_dir))
from gfpgan import GFPGANer
@ -313,7 +322,7 @@ def load_GFPGAN():
instance = GFPGANer(model_path=model_path, upscale=1, arch='clean', channel_multiplier=2, bg_upsampler=None, device=torch.device(f'cuda:{opt.gpu}'))
return instance
def load_RealESRGAN(model_name: str):
def load_RealESRGAN(model_name: str, checking = False):
from basicsr.archs.rrdbnet_arch import RRDBNet
RealESRGAN_models = {
'RealESRGAN_x4plus': RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4),
@ -323,7 +332,8 @@ def load_RealESRGAN(model_name: str):
model_path = os.path.join(RealESRGAN_dir, 'experiments/pretrained_models', model_name + '.pth')
if not os.path.isfile(model_path):
raise Exception(model_name+".pth not found at path "+model_path)
if checking == True:
return True
sys.path.append(os.path.abspath(RealESRGAN_dir))
from realesrgan import RealESRGANer
@ -341,32 +351,38 @@ def load_RealESRGAN(model_name: str):
GFPGAN = None
if os.path.exists(GFPGAN_dir):
try:
GFPGAN = load_GFPGAN()
print("Loaded GFPGAN")
GFPGAN = load_GFPGAN(checking=True)
print("Found GFPGAN")
except Exception:
import traceback
print("Error loading GFPGAN:", file=sys.stderr)
print(traceback.format_exc(), file=sys.stderr)
RealESRGAN = None
def try_loading_RealESRGAN(model_name: str):
def try_loading_RealESRGAN(model_name: str,checking=False):
global RealESRGAN
if os.path.exists(RealESRGAN_dir):
try:
RealESRGAN = load_RealESRGAN(model_name) # TODO: Should try to load both models before giving up
RealESRGAN = load_RealESRGAN(model_name,checking) # TODO: Should try to load both models before giving up
if checking == True:
print("Found RealESRGAN")
return True
print("Loaded RealESRGAN with model "+RealESRGAN.model.name)
except Exception:
import traceback
print("Error loading RealESRGAN:", file=sys.stderr)
print(traceback.format_exc(), file=sys.stderr)
try_loading_RealESRGAN('RealESRGAN_x4plus')
try_loading_RealESRGAN('RealESRGAN_x4plus',checking=True)
LDSR = None
def try_loading_LDSR(model_name: str):
def try_loading_LDSR(model_name: str,checking=False):
global LDSR
if os.path.exists(LDSR_dir):
try:
LDSR = load_LDSR() # TODO: Should try to load both models before giving up
LDSR = load_LDSR(checking=True) # TODO: Should try to load both models before giving up
if checking == True:
print("Found LDSR")
return True
print("Latent Diffusion Super Sampling (LDSR) model loaded")
except Exception:
import traceback
@ -374,7 +390,7 @@ def try_loading_LDSR(model_name: str):
print(traceback.format_exc(), file=sys.stderr)
else:
print("LDSR not found at path, please make sure you have cloned the LDSR repo to ./src/latent-diffusion/")
try_loading_LDSR('model')
try_loading_LDSR('model',checking=True)
def load_SD_model():
if opt.optimized:
@ -576,7 +592,7 @@ def check_prompt_length(prompt, comments):
def save_sample(image, sample_path_i, filename, jpg_sample, prompts, seeds, width, height, steps, cfg_scale,
normalize_prompt_weights, use_GFPGAN, write_info_files, write_sample_info_to_log_file, prompt_matrix, init_img, uses_loopback, uses_random_seed_loopback, skip_save,
skip_grid, sort_samples, sampler_name, ddim_eta, n_iter, batch_size, i, denoising_strength, resize_mode, skip_metadata):
skip_grid, sort_samples, sampler_name, ddim_eta, n_iter, batch_size, i, denoising_strength, resize_mode, skip_metadata=False):
filename_i = os.path.join(sample_path_i, filename)
if not jpg_sample:
if opt.save_metadata and not skip_metadata:
@ -764,7 +780,7 @@ def process_images(
fp, ddim_eta=0.0, do_not_save_grid=False, normalize_prompt_weights=True, init_img=None, init_mask=None,
keep_mask=False, mask_blur_strength=3, denoising_strength=0.75, resize_mode=None, uses_loopback=False,
uses_random_seed_loopback=False, sort_samples=True, write_info_files=True, write_sample_info_to_log_file=False, jpg_sample=False,
variant_amount=0.0, variant_seed=None,imgProcessorTask=True, job_info: JobInfo = None):
variant_amount=0.0, variant_seed=None,imgProcessorTask=False, job_info: JobInfo = None):
"""this is the main loop that both txt2img and img2img use; it calls func_init once inside all the scopes and func_sample once per batch"""
assert prompt is not None
torch_gc()
@ -944,15 +960,13 @@ def process_images(
save_sample(gfpgan_image, sample_path_i, gfpgan_filename, jpg_sample, prompts, seeds, width, height, steps, cfg_scale,
normalize_prompt_weights, use_GFPGAN, write_info_files, write_sample_info_to_log_file, prompt_matrix, init_img, uses_loopback, uses_random_seed_loopback, skip_save,
skip_grid, sort_samples, sampler_name, ddim_eta, n_iter, batch_size, i, denoising_strength, resize_mode, False)
#output_images.append(gfpgan_image) #287
output_images.append(gfpgan_image) #287
#if simple_templating:
# grid_captions.append( captions[i] + "\ngfpgan" )
if use_RealESRGAN and RealESRGAN is not None and not use_GFPGAN:
skip_save = True # #287 >_>
torch_gc()
if RealESRGAN.model.name != realesrgan_model_name:
try_loading_RealESRGAN(realesrgan_model_name)
output, img_mode = RealESRGAN.enhance(original_sample[:,:,::-1])
esrgan_filename = original_filename + '-esrgan4x'
esrgan_sample = output[:,:,::-1]
@ -960,7 +974,7 @@ skip_grid, sort_samples, sampler_name, ddim_eta, n_iter, batch_size, i, denoisin
save_sample(esrgan_image, sample_path_i, esrgan_filename, jpg_sample, prompts, seeds, width, height, steps, cfg_scale,
normalize_prompt_weights, use_GFPGAN,write_info_files, write_sample_info_to_log_file, prompt_matrix, init_img, uses_loopback, uses_random_seed_loopback, skip_save,
skip_grid, sort_samples, sampler_name, ddim_eta, n_iter, batch_size, i, denoising_strength, resize_mode, False)
#output_images.append(esrgan_image) #287
output_images.append(esrgan_image) #287
#if simple_templating:
# grid_captions.append( captions[i] + "\nesrgan" )
@ -969,8 +983,6 @@ skip_grid, sort_samples, sampler_name, ddim_eta, n_iter, batch_size, i, denoisin
torch_gc()
cropped_faces, restored_faces, restored_img = GFPGAN.enhance(x_sample[:,:,::-1], has_aligned=False, only_center_face=False, paste_back=True)
gfpgan_sample = restored_img[:,:,::-1]
if RealESRGAN.model.name != realesrgan_model_name:
try_loading_RealESRGAN(realesrgan_model_name)
output, img_mode = RealESRGAN.enhance(gfpgan_sample[:,:,::-1])
gfpgan_esrgan_filename = original_filename + '-gfpgan-esrgan4x'
gfpgan_esrgan_sample = output[:,:,::-1]
@ -978,12 +990,13 @@ skip_grid, sort_samples, sampler_name, ddim_eta, n_iter, batch_size, i, denoisin
save_sample(gfpgan_esrgan_image, sample_path_i, gfpgan_esrgan_filename, jpg_sample, prompts, seeds, width, height, steps, cfg_scale,
normalize_prompt_weights, use_GFPGAN, write_info_files, write_sample_info_to_log_file, prompt_matrix, init_img, uses_loopback, uses_random_seed_loopback, skip_save,
skip_grid, sort_samples, sampler_name, ddim_eta, n_iter, batch_size, i, denoising_strength, resize_mode, False)
#output_images.append(gfpgan_esrgan_image) #287
output_images.append(gfpgan_esrgan_image) #287
#if simple_templating:
# grid_captions.append( captions[i] + "\ngfpgan_esrgan" )
#if imgProcessorTask == True:
# output_images.append(image)
# this flag is used for imgProcessorTasks like GoBig, will return the image without saving it
if imgProcessorTask == True:
output_images.append(image)
if not skip_save:
save_sample(image, sample_path_i, filename, jpg_sample, prompts, seeds, width, height, steps, cfg_scale,
@ -1060,7 +1073,6 @@ def txt2img(prompt: str, ddim_steps: int, sampler_name: str, toggles: List[int],
outpath = opt.outdir_txt2img or opt.outdir or "outputs/txt2img-samples"
err = False
seed = seed_to_int(seed)
prompt_matrix = 0 in toggles
normalize_prompt_weights = 1 in toggles
skip_save = 2 not in toggles
@ -1071,7 +1083,15 @@ def txt2img(prompt: str, ddim_steps: int, sampler_name: str, toggles: List[int],
jpg_sample = 7 in toggles
use_GFPGAN = 8 in toggles
use_RealESRGAN = 9 in toggles
ModelLoader(['model'],True,False)
if use_GFPGAN and not use_RealESRGAN:
ModelLoader(['GFPGAN'],True,False)
ModelLoader(['RealESRGAN'],False,True)
if use_RealESRGAN and not use_GFPGAN:
ModelLoader(['GFPGAN'],False,True)
ModelLoader(['RealESRGAN'],True,False,realesrgan_model_name)
if use_RealESRGAN and use_GFPGAN:
ModelLoader(['GFPGAN','RealESRGAN'],True,False,realesrgan_model_name)
if sampler_name == 'PLMS':
sampler = PLMSSampler(model)
elif sampler_name == 'DDIM':
@ -1206,7 +1226,15 @@ def img2img(prompt: str, image_editor_mode: str, init_info: Dict[str,Image.Image
jpg_sample = 9 in toggles
use_GFPGAN = 10 in toggles
use_RealESRGAN = 11 in toggles
ModelLoader(['model'],True,False)
if use_GFPGAN and not use_RealESRGAN:
ModelLoader(['GFPGAN'],True,False)
ModelLoader(['RealESRGAN'],False,True)
if use_RealESRGAN and not use_GFPGAN:
ModelLoader(['GFPGAN'],False,True)
ModelLoader(['RealESRGAN'],True,False,realesrgan_model_name)
if use_RealESRGAN and use_GFPGAN:
ModelLoader(['GFPGAN','RealESRGAN'],True,False,realesrgan_model_name)
if sampler_name == 'DDIM':
sampler = DDIMSampler(model)
elif sampler_name == 'k_dpm_2_a':
@ -1489,7 +1517,8 @@ def slerp(device, t, v0:torch.Tensor, v1:torch.Tensor, DOT_THRESHOLD=0.9995):
def imgproc(image,image_batch,imgproc_prompt,imgproc_toggles, imgproc_upscale_toggles,imgproc_realesrgan_model_name,imgproc_sampling, imgproc_steps, imgproc_height, imgproc_width, imgproc_cfg, imgproc_denoising, imgproc_seed,imgproc_gfpgan_strength):
def imgproc(image,image_batch,imgproc_prompt,imgproc_toggles, imgproc_upscale_toggles,imgproc_realesrgan_model_name,imgproc_sampling,
imgproc_steps, imgproc_height, imgproc_width, imgproc_cfg, imgproc_denoising, imgproc_seed,imgproc_gfpgan_strength,imgproc_ldsr_steps,imgproc_ldsr_pre_downSample,imgproc_ldsr_post_downSample):
outpath = opt.outdir_imglab or opt.outdir or "outputs/imglab-samples"
output = []
@ -1734,8 +1763,9 @@ def imgproc(image,image_batch,imgproc_prompt,imgproc_toggles, imgproc_upscale_to
torch.cuda.empty_cache()
return combined_image
def processLDSR(image):
result = LDSR.superResolution(image)
return result
result = LDSR.superResolution(image,int(imgproc_ldsr_steps),str(imgproc_ldsr_pre_downSample),str(imgproc_ldsr_post_downSample))
return result
if image_batch != None:
if image != None:
@ -1755,9 +1785,27 @@ def imgproc(image,image_batch,imgproc_prompt,imgproc_toggles, imgproc_upscale_to
if len(images) > 0:
print("Processing images...")
#pre load models not in loop
if 0 in imgproc_toggles:
ModelLoader(['RealESGAN','LDSR'],False,True) # Unload unused models
ModelLoader(['GFPGAN'],True,False) # Load used models
if 1 in imgproc_toggles:
if imgproc_upscale_toggles == 0:
ModelLoader(['GFPGAN','LDSR'],False,True) # Unload unused models
ModelLoader(['RealESGAN'],True,False,imgproc_realesrgan_model_name) # Load used models
elif imgproc_upscale_toggles == 1:
ModelLoader(['GFPGAN','LDSR'],False,True) # Unload unused models
ModelLoader(['RealESGAN','model'],True,False) # Load used models
elif imgproc_upscale_toggles == 2:
ModelLoader(['model','GFPGAN','RealESGAN'],False,True) # Unload unused models
ModelLoader(['LDSR'],True,False) # Load used models
elif imgproc_upscale_toggles == 3:
ModelLoader(['GFPGAN','LDSR'],False,True) # Unload unused models
ModelLoader(['RealESGAN','model'],True,False,imgproc_realesrgan_model_name) # Load used models
for image in images:
if 0 in imgproc_toggles:
ModelLoader(['model','RealESGAN','LDSR'],False,True) # Unload unused models
#recheck if GFPGAN is loaded since it's the only model that can be loaded in the loop as well
ModelLoader(['GFPGAN'],True,False) # Load used models
image = processGFPGAN(image,imgproc_gfpgan_strength)
outpathDir = os.path.join(outpath,'GFPGAN')
@ -1770,9 +1818,6 @@ def imgproc(image,image_batch,imgproc_prompt,imgproc_toggles, imgproc_upscale_to
save_sample(image, outpathDir, outFilename, False, None, None, None, None, None, None, None, None, None, None, None, None, None, False, None, None, None, None, None, None, None, None, None, True)
if 1 in imgproc_toggles:
if imgproc_upscale_toggles == 0:
ModelLoader(['model','GFPGAN','LDSR'],False,True) # Unload unused models
ModelLoader(['RealESGAN'],True,False) # Load used models
image = processRealESRGAN(image)
outpathDir = os.path.join(outpath,'RealESRGAN')
os.makedirs(outpathDir, exist_ok=True)
@ -1782,9 +1827,6 @@ def imgproc(image,image_batch,imgproc_prompt,imgproc_toggles, imgproc_upscale_to
save_sample(image, outpathDir, outFilename, False, None, None, None, None, None, None, None, None, None, None, None, None, None, False, None, None, None, None, None, None, None, None, None, True)
elif imgproc_upscale_toggles == 1:
ModelLoader(['GFPGAN','LDSR'],False,True) # Unload unused models
ModelLoader(['RealESGAN','model'],True,False) # Load used models
image = processGoBig(image)
outpathDir = os.path.join(outpath,'GoBig')
os.makedirs(outpathDir, exist_ok=True)
@ -1794,9 +1836,6 @@ def imgproc(image,image_batch,imgproc_prompt,imgproc_toggles, imgproc_upscale_to
save_sample(image, outpathDir, outFilename, False, None, None, None, None, None, None, None, None, None, None, None, None, None, False, None, None, None, None, None, None, None, None, None, True)
elif imgproc_upscale_toggles == 2:
ModelLoader(['model','GFPGAN','RealESGAN'],False,True) # Unload unused models
ModelLoader(['LDSR'],True,False) # Load used models
image = processLDSR(image)
outpathDir = os.path.join(outpath,'LDSR')
os.makedirs(outpathDir, exist_ok=True)
@ -1806,9 +1845,6 @@ def imgproc(image,image_batch,imgproc_prompt,imgproc_toggles, imgproc_upscale_to
save_sample(image, outpathDir, outFilename, False, None, None, None, None, None, None, None, None, None, None, None, None, None, False, None, None, None, None, None, None, None, None, None, True)
elif imgproc_upscale_toggles == 3:
ModelLoader(['GFPGAN','LDSR'],False,True) # Unload unused models
ModelLoader(['RealESGAN','model'],True,False) # Load used models
image = processGoBig(image)
ModelLoader(['model','GFPGAN','RealESGAN'],False,True) # Unload unused models
ModelLoader(['LDSR'],True,False) # Load used models
@ -1822,13 +1858,13 @@ def imgproc(image,image_batch,imgproc_prompt,imgproc_toggles, imgproc_upscale_to
save_sample(image, outpathDir, outFilename, None, None, None, None, None, None, None, None, None, None, None, None, None, None, False, None, None, None, None, None, None, None, None, None, True)
#LDSR is always unloaded to avoid memory issues
ModelLoader(['LDSR'],False,True)
print("Reloading default models...")
ModelLoader(['model','RealESGAN','GFPGAN'],True,False) # load back models
#ModelLoader(['LDSR'],False,True)
#print("Reloading default models...")
#ModelLoader(['model','RealESGAN','GFPGAN'],True,False) # load back models
print("Done.")
return output
def ModelLoader(models,load=False,unload=False):
def ModelLoader(models,load=False,unload=False,imgproc_realesrgan_model_name='RealESRGAN_x4plus'):
#get global variables
global_vars = globals()
#check if m is in globals
@ -1846,7 +1882,7 @@ def ModelLoader(models,load=False,unload=False):
print('Unloaded ' + m)
if load:
for m in models:
if m not in global_vars:
if m not in global_vars or m in global_vars and type(global_vars[m]) == bool:
#if it isn't, load it
if m == 'GFPGAN':
global_vars[m] = load_GFPGAN()
@ -1857,17 +1893,18 @@ def ModelLoader(models,load=False,unload=False):
global_vars[m+'CS'] = sdLoader[1]
global_vars[m+'FS'] = sdLoader[2]
elif m == 'RealESRGAN':
global_vars[m] = load_RealESRGAN('RealESRGAN_x4plus')
global_vars[m] = load_RealESRGAN(imgproc_realesrgan_model_name)
elif m == 'LDSR':
global_vars[m] = load_LDSR()
if m =='model':
m='Stable Diffusion'
print('Loaded ' + m)
torch_gc()
def run_GFPGAN(image, strength):
ModelLoader(['LDSR','RealESRGAN'],False,True)
ModelLoader(['GFPGAN'],True,False)
image = image.convert("RGB")
cropped_faces, restored_faces, restored_img = GFPGAN.enhance(np.array(image, dtype=np.uint8), has_aligned=False, only_center_face=False, paste_back=True)
@ -1879,6 +1916,8 @@ def run_GFPGAN(image, strength):
return res
def run_RealESRGAN(image, model_name: str):
ModelLoader(['GFPGAN','LDSR'],False,True)
ModelLoader(['RealESRGAN'],True,False)
if RealESRGAN.model.name != model_name:
try_loading_RealESRGAN(model_name)
@ -1972,13 +2011,12 @@ img2img_toggles = [
'Write sample info to one file',
'jpg samples',
]
"""
# removed for now becuase of Image Lab implementation
if GFPGAN is not None:
img2img_toggles.append('Fix faces using GFPGAN')
if RealESRGAN is not None:
img2img_toggles.append('Upscale images using RealESRGAN')
"""
img2img_mask_modes = [
"Keep masked area",
"Regenerate only masked area",