Please download GFPGAN to activate face fixing features, instructions are available at the Github
-Please download GFPGAN to activate face fixing features, instructions are available at the Github
+Please download LDSR to activate more upscale features, instructions are available at the Github
-Please download LDSR to activate more upscale features, instructions are available at the Github
+Please download RealESRGAN to activate upscale features, instructions are available at the Github
""") - + 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 = { diff --git a/frontend/ui_functions.py b/frontend/ui_functions.py index 633ae1a..fa731e0 100644 --- a/frontend/ui_functions.py +++ b/frontend/ui_functions.py @@ -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)) diff --git a/webui.py b/webui.py index af0f4e5..bef9b45 100644 --- a/webui.py +++ b/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",