mirror of
https://github.com/sd-webui/stable-diffusion-webui.git
synced 2024-12-14 06:35:14 +03:00
[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
GitHub
parent
b61b453a02
commit
de95a09593
13
README.md
13
README.md
@ -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.
|
||||
|
||||
|
||||
|
||||
|
||||
### Sampling method selection
|
||||
|
||||
@ -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;
|
||||
|
||||
@ -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 = {
|
||||
|
||||
@ -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
148
webui.py
@ -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",
|
||||
|
||||
Loading…
Reference in New Issue
Block a user