├── latents └── Put saved latents here.txt ├── run.bat ├── Images └── Output images will be saved here.txt ├── additional_requirements.txt ├── faces └── Detected faces are extracted here.txt ├── models ├── put downloaded main model here.txt └── CommonModels │ ├── VaeDecoder │ └── put downloaded model here.txt │ ├── VaeEncoder │ └── put downloaded model here.txt │ └── Text-Encoder-ClipSkip1 │ └── put downloaded model here.txt ├── activate.bat ├── deactivate.bat ├── modules ├── database │ └── Database will be created here.txt ├── dis-isnet │ └── put isnet model here, download from huggingface.txt ├── modules.py ├── wildcards_module.py ├── styles_module.py ├── image_to_numpy_module.py └── latent_formula_processor.py ├── Scripts ├── wildcards │ ├── Descompress all txt files(or only needed) from zip file.txt │ └── wilcards_text_files.zip ├── wildcards.py └── utils.py ├── Engine ├── config_files │ ├── TextEnc_Preferences.json │ ├── VAE_Preferences.json │ ├── UIConfig.json │ ├── Styles.json │ └── ControlNet.json ├── __init__.py ├── Pipelines │ ├── inpaint_pipeline.py │ ├── img2img_pipeline.py │ ├── Controlnet_pipeline.py │ └── Vae_and_Text_Encoders.py ├── engine_common_funcs.py ├── General_parameters.py └── txt2img_pipe_sub.py ├── requirements.txt ├── UI ├── tools_ui │ ├── danbooru_ui.py │ ├── clipskip_ui.py │ └── clipskip_config_ui.py ├── config_ui_wildcards.py ├── UI_placement_areas.py ├── UI_common_funcs.py ├── styles_ui.py ├── config_ui_general.py ├── edit_styles_ui.py ├── config_ui_ControlNet.py ├── config_ui_Vae.py ├── instructp2p_ui.py ├── config_ui_TextEncoder.py ├── Inpaint_ui.py ├── ControlNet_ui.py ├── config_ui_engine.py ├── Img2Img_ui.py └── ui_face_tools.py ├── latent_to_pil.py ├── ONNX-StableDiffusion.py ├── README.md └── IMG-Strip.py /latents/Put saved latents here.txt: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /run.bat: -------------------------------------------------------------------------------- 1 | py -O ONNX-StableDiffusion.py -------------------------------------------------------------------------------- /Images/Output images will be saved here.txt: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /additional_requirements.txt: -------------------------------------------------------------------------------- 1 | insightface 2 | -------------------------------------------------------------------------------- /faces/Detected faces are extracted here.txt: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /models/put downloaded main model here.txt: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /activate.bat: -------------------------------------------------------------------------------- 1 | ./venv/scripts/activate.bat 2 | -------------------------------------------------------------------------------- /deactivate.bat: -------------------------------------------------------------------------------- 1 | ./venv/scripts/deactivate.bat 2 | -------------------------------------------------------------------------------- /modules/database/Database will be created here.txt: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /models/CommonModels/VaeDecoder/put downloaded model here.txt: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /models/CommonModels/VaeEncoder/put downloaded model here.txt: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /models/CommonModels/Text-Encoder-ClipSkip1/put downloaded model here.txt: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /Scripts/wildcards/Descompress all txt files(or only needed) from zip file.txt: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /Engine/config_files/TextEnc_Preferences.json: -------------------------------------------------------------------------------- 1 | ["./models/CommonModels/Text-Encoder-ClipSkip1", "model", "", "", ""] -------------------------------------------------------------------------------- /Engine/config_files/VAE_Preferences.json: -------------------------------------------------------------------------------- 1 | ["./models/CommonModels/VaeDecoder", "model", "", "./models/CommonModels/VaeEncoder", "model", ""] -------------------------------------------------------------------------------- /modules/dis-isnet/put isnet model here, download from huggingface.txt: -------------------------------------------------------------------------------- 1 | https://huggingface.co/Neus/CommonModels/tree/main/dis-isnet 2 | -------------------------------------------------------------------------------- /Scripts/wildcards/wilcards_text_files.zip: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/NeusZimmer/ONNX-ModularUI-StableDiffusion/HEAD/Scripts/wildcards/wilcards_text_files.zip -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | protobuf < 4.0 2 | torch 3 | numpy 4 | transformers 5 | diffusers 6 | ftfy 7 | spacy 8 | scipy 9 | safetensors 10 | gradio == 3.50.2 11 | omegaconf 12 | onnx 13 | onnxconverter-common 14 | onnxruntime-directml 15 | opencv-python 16 | -------------------------------------------------------------------------------- /Engine/config_files/UIConfig.json: -------------------------------------------------------------------------------- 1 | {"models_dir": "./models", "output_path": "./Images", "Txt2img_Tab": true, "InPaint_Tab": true, "Img2Img_Tab": true, "InstructP2P_Tab": 0, "ControlNet_Tab": true, "Tools_Tab": false, "Advanced_Config": true, "GradioPort": "7860", "_UI_Configuration__loaded": true} -------------------------------------------------------------------------------- /Engine/__init__.py: -------------------------------------------------------------------------------- 1 | from .Pipelines.Vae_and_Text_Encoders import Vae_and_Text_Encoders 2 | from .Pipelines.schedulers import SchedulersConfig 3 | from .Pipelines.txt2img_pipeline import txt2img_pipe 4 | from .Pipelines.txt2img_hires import txt2img_hires_pipe 5 | from .Pipelines.img2img_pipeline import img2img_pipe 6 | from .Pipelines.Controlnet_pipeline import ControlNet_pipe 7 | from .Pipelines.inpaint_pipeline import inpaint_pipe 8 | 9 | 10 | -------------------------------------------------------------------------------- /Engine/config_files/Styles.json: -------------------------------------------------------------------------------- 1 | {"None": true, "Abstract": " Abstract, 8K, photo shot | , in a modern and abstract setting, with bold and colorful abstract art, blurred background, bright lighting,", "Hatching": " | (Contoured hatching)", "Photo": " (RAW Photo)| , (high quality, highres), raw photo", "HQ_Anime": "Anime style, masterpiece, best quality,absurdres, Hyperdetailed, masterpiece, best quality, | ,HDR (High Dynamic Range)", "Cartoon": "Still of an animated movie, | ,Style of classic cartoon movies", "LineArt": "(cartoon) |, Ink drawing line art, ((matte painting fantasy concept art))"} -------------------------------------------------------------------------------- /Engine/config_files/ControlNet.json: -------------------------------------------------------------------------------- 1 | {"canny_active": true, "canny_path": "D:\\onnx-win\\files\\models\\control\\canny", "depth_active": false, "depth_path": "D:\\onnx-win\\files\\models\\control\\depth", "hed_active": false, "hed_path": "D:\\onnx-win\\files\\models\\control\\hed", "mlsd_active": false, "mlsd_path": "D:\\onnx-win\\files\\models\\control\\mlsd", "normal_active": false, "normal_path": "D:\\onnx-win\\files\\models\\control\\normal", "openpose_active": false, "openpose_path": "D:\\onnx-win\\files\\models\\control\\openpose", "seg_active": false, "seg_path": "C:\\AMD_ML\\Conversion\\controlnet\\onnx_stable_diffusion_controlnet\\model\\control_v11f1e_sd15_tile\\controlnet"} -------------------------------------------------------------------------------- /UI/tools_ui/danbooru_ui.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | def show_danbooru_area(): 6 | global image_in 7 | with gr.Row(): 8 | with gr.Column(variant="compact"): 9 | image_in = gr.Image(label="input image", type="pil", elem_id="image_init") 10 | with gr.Row(): 11 | apply_btn = gr.Button("Analyze image with Deep DanBooru", variant="primary") 12 | mem_btn = gr.Button("Unload from memory") 13 | with gr.Row(): 14 | results = gr.Textbox(value="", lines=8, label="Results") 15 | 16 | mem_btn.click(fn=unload_DanBooru, inputs=results , outputs=results) 17 | apply_btn.click(fn=analyze_DanBooru, inputs=image_in , outputs=results) -------------------------------------------------------------------------------- /UI/config_ui_wildcards.py: -------------------------------------------------------------------------------- 1 | import gradio as gr 2 | from Engine import General_parameters as params 3 | 4 | global debug 5 | debug = True 6 | 7 | 8 | def show_wilcards_configuration(): 9 | with gr.Blocks(title="ONNX Difussers UI-2") as test2: 10 | with gr.Accordion(label="Wilcards options",open=False): 11 | gr.Markdown("Still not implemented: Wildcards not recursive & Not working option to change marker of wildcards") 12 | with gr.Row(): 13 | Wilcards_Select=gr.Radio(["Yes","No","Yes,not recursive"], label="Wilcards options", info="Activate Yes/No",value="Yes") 14 | Wilcards_Select.change(fn=Generic_Select_Option,inputs=Wilcards_Select,outputs=None) 15 | 16 | def Generic_Select_Option(Radio_Select): 17 | config=params.UI_Configuration() 18 | if Radio_Select == "Yes": 19 | config.wildcards_activated=True 20 | print(params.UI_Configuration().wildcards_activated) 21 | else: 22 | config.wildcards_activated=False 23 | print(params.UI_Configuration().wildcards_activated) 24 | 25 | -------------------------------------------------------------------------------- /UI/UI_placement_areas.py: -------------------------------------------------------------------------------- 1 | #This file exist mainly for better area definition understanding. 2 | 3 | def show_new_tabs(list_modules,tab="main_ui"): 4 | pass 5 | 6 | def show_tools_area(list_modules): 7 | _all_areas_process(list_modules,"tools_area") #Area of tool modules 8 | 9 | def show_prompt_preprocess_area(list_modules): 10 | _all_areas_process(list_modules,"prompt_process") #Area of modules for processing prompts 11 | 12 | def show_image_postprocess_area(list_modules): 13 | _all_areas_process(list_modules,"image_postprocess")#Area of modules for postprocessing images 14 | 15 | def show_footer_area(list_modules,tab="NoUI"): 16 | _all_areas_process(list_modules,"footer",tab)#Area of modules for postprocessing images 17 | 18 | def _all_areas_process(list_modules,area_name,tab="NoUI"): 19 | for module in list_modules: 20 | show=False 21 | if tab=="main_ui" and module['is_global_ui'](): show=True 22 | if tab!="main_ui" and not module['is_global_ui'](): 23 | if module['ui_position']==area_name: 24 | show= True 25 | 26 | if show: 27 | module['show']() #modules[2]= show, #modules[3] =process 28 | 29 | 30 | if __name__ == "main": 31 | print("This file is not intended to run as standalone") 32 | -------------------------------------------------------------------------------- /latent_to_pil.py: -------------------------------------------------------------------------------- 1 | from Engine import Vae_and_Text_Encoders,engine_common_funcs 2 | from optimum.onnxruntime.modeling_ort import ORTModel 3 | import numpy as np 4 | import os,sys 5 | latent_path="./latents" 6 | latent_list = [] 7 | divider=1 8 | 9 | for i, arg in enumerate(sys.argv): 10 | print(f"Argument {i:>6}: {arg}") 11 | 12 | try: 13 | with os.scandir(latent_path) as scan_it: 14 | for entry in scan_it: 15 | if ".npy" in entry.name: 16 | latent_list.append(entry.name) 17 | print(latent_list) 18 | except: 19 | print("Not numpys found.Wrong Directory or no files inside?") 20 | 21 | vae_path="D:\\models\\ModelosVAE\\vae_decoder-standar" 22 | vaedec_sess = ORTModel.load_model(vae_path+"/model.onnx", provider="DmlExecutionProvider", provider_options={'device_id': 1}) 23 | 24 | for latent in latent_list: 25 | try: 26 | loaded_latent=np.load(f"./latents/{latent}") 27 | loaded_latent = 1 / 0.18215 * loaded_latent 28 | 29 | import torch.nn.functional as F 30 | import torch.onnx 31 | import torch 32 | latent1=torch.from_numpy(loaded_latent) 33 | print(latent1.size()) 34 | latent1= F.interpolate(latent1,size=(int(latent1.size()[2]/divider), int(latent1.size()[3]/divider)), mode='bilinear') 35 | print(latent1.size()) 36 | loaded_latent = latent1.numpy() 37 | 38 | image= vaedec_sess.run(['sample'],{'latent_sample': loaded_latent})[0] 39 | name= latent[:-3] 40 | name= name+"png" 41 | image = np.clip(image / 2 + 0.5, 0, 1) 42 | image = image.transpose((0, 2, 3, 1)) 43 | image = engine_common_funcs.numpy_to_pil(image)[0] 44 | image.save(f"./latents/{name}",optimize=True) 45 | print(f"Saved:{name}") 46 | except: 47 | print(f"Error opening/processing:{latent}") 48 | 49 | del vaedec_sess -------------------------------------------------------------------------------- /UI/UI_common_funcs.py: -------------------------------------------------------------------------------- 1 | import gc,os 2 | from Engine import pipelines_engines 3 | 4 | from Engine.General_parameters import running_config 5 | from Engine.General_parameters import UI_Configuration 6 | #from Engine.Pipelines.txt2img_hires import txt2img_hires_pipe 7 | #from Engine.Pipelines.txt2img_pipeline import txt2img_pipe 8 | from Engine import ( 9 | txt2img_pipe, 10 | Vae_and_Text_Encoders, 11 | txt2img_hires_pipe, 12 | ControlNet_pipe, 13 | img2img_pipe 14 | ) 15 | 16 | 17 | def clean_memory_click(): 18 | print("Cleaning Memory") 19 | Vae_and_Text_Encoders().unload_from_memory() 20 | txt2img_pipe().unload_from_memory() 21 | txt2img_hires_pipe().unload_from_memory() 22 | ControlNet_pipe().unload_from_memory() 23 | img2img_pipe().unload_from_memory() 24 | #pipelines_engines.inpaint_pipe().unload_from_memory() 25 | #pipelines_engines.instruct_p2p_pipe().unload_from_memory() 26 | 27 | gc.collect() 28 | 29 | def cancel_iteration(): 30 | running_config().Running_information.update({"cancelled":True}) 31 | print("\nCancelling at the end of the current iteration") 32 | 33 | 34 | 35 | def get_model_list(pipeline=None): 36 | model_list = [] 37 | inpaint_model_list = [] 38 | controlnet_model_list = [] 39 | i2p_model_list = [] 40 | txt2img_model_list = [] 41 | 42 | try: 43 | with os.scandir(UI_Configuration().models_dir) as scan_it: 44 | for entry in scan_it: 45 | if entry.is_dir(): 46 | model_list.append(entry.name) 47 | 48 | except: 49 | model_list.append("Models directory does not exist, configure it") 50 | 51 | if len(model_list)>0: 52 | for model in model_list: 53 | if "inpaint" in model.lower(): inpaint_model_list.append(model) 54 | elif "controlnet" in model.lower(): controlnet_model_list.append(model) 55 | elif "ip2p" in model.lower(): i2p_model_list.append(model) 56 | else: txt2img_model_list.append(model) 57 | 58 | if pipeline=="txt2img": retorno=txt2img_model_list 59 | elif pipeline=="inpaint": retorno=inpaint_model_list 60 | elif pipeline=="controlnet": retorno= controlnet_model_list 61 | elif pipeline=="ip2p": retorno= i2p_model_list 62 | else: retorno= model_list 63 | 64 | return retorno -------------------------------------------------------------------------------- /Scripts/wildcards.py: -------------------------------------------------------------------------------- 1 | # Wildcards 2 | #V3 Extended to provide info on status area and recursive wildcards (including one wildcard into a line of a wildcard file for increased variability) 3 | #example a prompt could be:"An incrincated drawing of __example__ " and the inside the example.txt wildcard file a line: __test__ could include the line "A __color__ __vehicle__ running trough __landscape__ " 4 | #An extension modificated version of a script from https://github.com/jtkelm2/stable-diffusion-webui-1/blob/main/scripts/wildcards.py 5 | #Idea originated , but modified, from module wildcards of AUTOMATIC111 6 | #Allows you to use `__name__` syntax in your prompt to get a random line from a file named `name.txt` in the wildcards directory. 7 | 8 | 9 | 10 | import os 11 | import random 12 | import sys 13 | 14 | warned_about_files = {} 15 | wildcard_dir = os.getcwd()+"\Scripts" 16 | #print(wildcard_dir) 17 | 18 | 19 | class WildcardsScript(): 20 | def title(self): 21 | return "Simple wildcards class for OnnxDiffusersUI" 22 | 23 | def replace_wildcard(self, text, gen): 24 | if " " in text or len(text) == 0: 25 | return text,False 26 | 27 | replacement_file = os.path.join(wildcard_dir, "wildcards", f"{text}.txt") 28 | if os.path.exists(replacement_file): 29 | with open(replacement_file, encoding="utf8") as f: 30 | changed_text=gen.choice(f.read().splitlines()) 31 | if "__" in changed_text: 32 | changed_text, not_used = self.process(changed_text) 33 | return changed_text,True 34 | else: 35 | if replacement_file not in warned_about_files: 36 | print(f"File {replacement_file} not found for the __{text}__ wildcard.", file=sys.stderr) 37 | warned_about_files[replacement_file] = 1 38 | 39 | return text,False 40 | 41 | def process(self, original_prompt): 42 | string_replaced="" 43 | new_prompt="" 44 | gen = random.Random() 45 | text_divisions=original_prompt.split("__") 46 | 47 | for chunk in text_divisions: 48 | text,changed=self.replace_wildcard(chunk, gen) 49 | if changed: 50 | string_replaced=string_replaced+"Wildcard:"+chunk+"-->"+text+"," 51 | new_prompt=new_prompt+text 52 | else: 53 | new_prompt=new_prompt+text 54 | 55 | return new_prompt, string_replaced 56 | -------------------------------------------------------------------------------- /UI/styles_ui.py: -------------------------------------------------------------------------------- 1 | import gradio as gr 2 | #from Engine.General_parameters import ControlNet_config 3 | from UI import styles_ui 4 | global styles_dict 5 | 6 | def show_styles_ui(): 7 | global styles_dict 8 | styles_dict= get_styles() 9 | styles_keys= list(styles_dict.keys()) 10 | if True: 11 | with gr.Accordion(label="Styles",open=False): 12 | gr.Markdown("Use your preferred Styles") 13 | with gr.Row(): 14 | Style_Select = gr.Radio(styles_keys,value=styles_keys[0],label="Available Styles") 15 | with gr.Row(): 16 | with gr.Accordion(label="Style - modificate running config",open=False): 17 | styletext_pre = gr.Textbox(value="", lines=2, label="Style previous text") 18 | styletext_post = gr.Textbox(value="", lines=2, label="Style posterior text") 19 | with gr.Row(): 20 | apply_btn = gr.Button("Apply Modified Style") 21 | reload_btn = gr.Button("Reload Styles") 22 | 23 | all_inputs=[Style_Select,styletext_pre,styletext_post] 24 | 25 | apply_btn.click(fn=saveto_memory_style, inputs=all_inputs, outputs=None) 26 | reload_btn.click(fn=reload_styles, inputs=None, outputs=Style_Select) 27 | Style_Select.change(fn=apply_styles, inputs=Style_Select, outputs=[styletext_pre,styletext_post]) 28 | 29 | def get_styles(): 30 | import json 31 | """dict={ 32 | "None":True, 33 | "StudioPhoto":"(RAW, 8k) |, studio lights,pseudo-impasto", 34 | "Style1":"(cartoon) |, Ink drawing line art", 35 | "Style2":"unity wallpaper, 8k, high quality, | masterpiece,(masterpiece, top quality, best quality)" 36 | }""" 37 | with open('./Engine/config_files/Styles.json', 'r') as openfile: 38 | jsonStr = json.load(openfile) 39 | 40 | jsonStr.update({"None":True}) 41 | return jsonStr 42 | 43 | def reload_styles(*args): 44 | global styles_dict 45 | styles_dict=get_styles() 46 | styles_keys= list(styles_dict.keys()) 47 | apply_styles("None") 48 | return gr.Radio.update(choices=styles_keys,value="None") 49 | 50 | 51 | def apply_styles(*args): 52 | global styles_dict 53 | dict=styles_dict 54 | style=args[0] 55 | 56 | from Engine.General_parameters import running_config 57 | Running_information= running_config().Running_information 58 | Running_information.update({"Style":False}) 59 | 60 | if style != "None": 61 | Running_information.update({"Style":dict[style]}) 62 | params=dict[style].split("|") 63 | return params[0],params[1] 64 | else: 65 | return "","" 66 | 67 | def saveto_memory_style(*args): 68 | import json 69 | global styles_dict 70 | styles_dict 71 | 72 | style=args[0] 73 | style_pre=args[1] 74 | style_post=args[2] 75 | 76 | 77 | from Engine.General_parameters import running_config 78 | Running_information= running_config().Running_information 79 | Running_information.update({"Style":False}) 80 | 81 | if style != "None": 82 | styles_dict.update({style:f"{style_pre}|{style_post}"}) 83 | Running_information.update({"Style":styles_dict[style]}) 84 | 85 | 86 | -------------------------------------------------------------------------------- /modules/modules.py: -------------------------------------------------------------------------------- 1 | """ 2 | Initialize the modules (__init__) to get the relevant info of what to do with them: 3 | 4 | The loaded modules must return a tuple of 4 elements 5 | 1st: a list of tab names where to be shown if ui component, 6 | 2nd: the area where to include the Ai and area to include the process to be done within the module 7 | 3rd and 4th: two funcions: one named "__call__" and another named "show" 8 | if no function, create one with the return being the same as the imput or a pass function 9 | 3rd: show: will be user to show the UI component of the module in the specified TAB & AREA, 10 | def show(*args): 11 | pass 12 | 4th: Call will be used to process a dictionay of parameters or a specific parameter, depending on where is to be included 13 | def __call__(*args): 14 | return args 15 | ATT: not implemented the area zones, fixed. 16 | tabs names:["txt2img","hires"] -- 17 | 18 | 19 | """ 20 | class Borg20: 21 | _shared_state = {} 22 | def __init__(self): 23 | self.__dict__ = self._shared_state 24 | 25 | 26 | 27 | class preProcess_modules(Borg20): 28 | all_modules=None 29 | _loaded=False 30 | 31 | 32 | def __init__(self): 33 | Borg20.__init__(self) 34 | if not self._loaded: 35 | self.__initclass__() 36 | 37 | def __str__(self): 38 | import json 39 | return json.dumps(self.__dict__) 40 | 41 | def __initclass__(self): 42 | #self.all_modules=self._launch_preprocess_modules() 43 | self.all_modules=self._load_all_modules() 44 | self._loaded=True 45 | 46 | def check_available_modules(self,tab_name): 47 | available_modules=[] 48 | for module in self.all_modules: 49 | if tab_name in module['tabs']: 50 | available_modules.append(module) 51 | 52 | return available_modules #do not use a self var, it will provide always the functions for the last UI tab loaded 53 | 54 | 55 | 56 | def _load_all_modules(*args,**kwargs): 57 | #def _load_preprocess_modules(*args,**kwargs): 58 | from importlib import import_module 59 | 60 | #lista=['library_module','wildcards_module','styles_module','image_to_numpy_module','reload_hires_module'] 61 | lista=['wildcards_module','styles_module','image_to_numpy_module'] 62 | modules_data=[] 63 | 64 | for elemento in lista: 65 | my_modulo=import_module('modules.'+elemento, package="StylesModule") 66 | modules_info=my_modulo.__init__("External Module %s Loaded" % elemento ) 67 | functions=(my_modulo.show,my_modulo.__call__,my_modulo.is_global_ui,my_modulo.is_global_function) 68 | modules_info.update({"show": functions[0]}) 69 | modules_info.update({"call": functions[1]}) 70 | modules_info.update({"is_global_ui": functions[2]}) 71 | modules_info.update({"is_global_function": functions[3]}) 72 | modules_data.append(modules_info) 73 | 74 | return modules_data # A list of dicts, one for each module 75 | 76 | 77 | 78 | if __name__ == "__main__": 79 | print("This is the module loader and is not intended to run as standalone") 80 | pass 81 | else: 82 | __name__ = "ExternalModuleLoader" 83 | -------------------------------------------------------------------------------- /UI/tools_ui/clipskip_ui.py: -------------------------------------------------------------------------------- 1 | import gradio as gr 2 | #from Engine.General_parameters import ControlNet_config 3 | from UI import styles_ui 4 | global styles_dict 5 | 6 | def show_styles_ui(): 7 | global styles_dict 8 | styles_dict= get_styles() 9 | styles_keys= list(styles_dict.keys()) 10 | if True: 11 | with gr.Accordion(label="Styles",open=False): 12 | gr.Markdown("Use your preferred Styles") 13 | with gr.Row(): 14 | Style_Select = gr.Radio(styles_keys,value=styles_keys[0],label="Available Styles") 15 | with gr.Row(): 16 | with gr.Accordion(label="Style - modificate running config",open=False): 17 | styletext_pre = gr.Textbox(value="", lines=2, label="Style previous text") 18 | styletext_post = gr.Textbox(value="", lines=2, label="Style posterior text") 19 | with gr.Row(): 20 | apply_btn = gr.Button("Apply Modified Style") 21 | reload_btn = gr.Button("Reload Styles") 22 | 23 | all_inputs=[Style_Select,styletext_pre,styletext_post] 24 | 25 | apply_btn.click(fn=saveto_memory_style, inputs=all_inputs, outputs=None) 26 | reload_btn.click(fn=reload_styles, inputs=None, outputs=Style_Select) 27 | Style_Select.change(fn=apply_styles, inputs=Style_Select, outputs=[styletext_pre,styletext_post]) 28 | 29 | def get_styles(): 30 | import json 31 | """dict={ 32 | "None":True, 33 | "StudioPhoto":"(RAW, 8k) |, studio lights,pseudo-impasto", 34 | "Style1":"(cartoon) |, Ink drawing line art", 35 | "Style2":"unity wallpaper, 8k, high quality, | masterpiece,(masterpiece, top quality, best quality)" 36 | }""" 37 | with open('./Engine/config_files/Styles.json', 'r') as openfile: 38 | jsonStr = json.load(openfile) 39 | 40 | jsonStr.update({"None":True}) 41 | return jsonStr 42 | 43 | def reload_styles(*args): 44 | global styles_dict 45 | styles_dict=get_styles() 46 | styles_keys= list(styles_dict.keys()) 47 | apply_styles("None") 48 | return gr.Radio.update(choices=styles_keys,value="None") 49 | 50 | 51 | def apply_styles(*args): 52 | global styles_dict 53 | dict=styles_dict 54 | style=args[0] 55 | 56 | from Engine.General_parameters import running_config 57 | Running_information= running_config().Running_information 58 | Running_information.update({"Style":False}) 59 | 60 | if style != "None": 61 | Running_information.update({"Style":dict[style]}) 62 | params=dict[style].split("|") 63 | return params[0],params[1] 64 | else: 65 | return "","" 66 | 67 | def saveto_memory_style(*args): 68 | import json 69 | global styles_dict 70 | styles_dict 71 | 72 | style=args[0] 73 | style_pre=args[1] 74 | style_post=args[2] 75 | 76 | 77 | from Engine.General_parameters import running_config 78 | Running_information= running_config().Running_information 79 | Running_information.update({"Style":False}) 80 | 81 | if style != "None": 82 | styles_dict.update({style:f"{style_pre}|{style_post}"}) 83 | Running_information.update({"Style":styles_dict[style]}) 84 | 85 | 86 | -------------------------------------------------------------------------------- /modules/wildcards_module.py: -------------------------------------------------------------------------------- 1 | """ 2 | Initialize: 3 | 4 | Return: dict of 3 elements 5 | 1st: a list of tab names where you want it to be included, 6 | 2nd: the area where to be shown in the UI 7 | 3rd: the point where to execute the call 8 | 9 | Additionally you need two funcions: one named "__call__" and another named "__show__" 10 | if no need for an specific function, create one with the return being the same as the imput or a pass function 11 | function __show__: the gradio elements to be shown into the UI component in the specified TAB & AREA, 12 | def show(*args): 13 | pass 14 | function __call__: this will be called to process a dictionay of parameters or a specific parameter, depending on where is to be included 15 | def __call__(*args): 16 | return args 17 | 18 | tabs names:["txt2img","hires"] -- 19 | 20 | 21 | """ 22 | 23 | 24 | 25 | def __init__(*args): 26 | __name__='WildcardModule' 27 | print(args[0]) 28 | #here a check of access of initializacion if needed 29 | #return (["txt2img","hires"],"prompt_process") 30 | 31 | return { 32 | 'tabs':["txt2img","hires"], 33 | 'ui_position':"prompt_process", 34 | 'func_processing':'prompt_process'} 35 | 36 | def __call__(*args): 37 | #what to do when the module is called 38 | #print("Processing wildcards Module") 39 | if args: 40 | datos=args[0] 41 | if type(datos)==dict: 42 | print("Dummy1-Wildcards") 43 | datos1=process(datos['prompt']) 44 | elif type(datos)==str: 45 | datos1=process(datos) 46 | #print("Dummy2") 47 | else: 48 | print("Not recognized input for module wildcards %s" % datos) 49 | return datos1 50 | 51 | def is_global_ui(): 52 | return False 53 | 54 | def is_global_function(): 55 | return False 56 | 57 | def show(): 58 | import gradio 59 | gradio.Markdown("Wildcards Module Activated") 60 | #pass 61 | 62 | 63 | def replace_wildcard(text, gen): 64 | import os,sys 65 | 66 | warned_about_files = {} 67 | wildcard_dir = os.getcwd()+"\Scripts" 68 | 69 | if " " in text or len(text) == 0: 70 | return text,False 71 | 72 | replacement_file = os.path.join(wildcard_dir, "wildcards", f"{text}.txt") 73 | if os.path.exists(replacement_file): 74 | with open(replacement_file, encoding="utf8") as f: 75 | changed_text=gen.choice(f.read().splitlines()) 76 | if "__" in changed_text: 77 | changed_text, not_used = self.process(changed_text) 78 | return changed_text,True 79 | else: 80 | if replacement_file not in warned_about_files: 81 | print(f"File {replacement_file} not found for the __{text}__ wildcard.", file=sys.stderr) 82 | warned_about_files[replacement_file] = 1 83 | 84 | return text,False 85 | 86 | def process(original_prompt): 87 | import random 88 | string_replaced="" 89 | new_prompt="" 90 | gen = random.Random() 91 | text_divisions=original_prompt.split("__") 92 | 93 | for chunk in text_divisions: 94 | text,changed=replace_wildcard(chunk, gen) 95 | if changed: 96 | string_replaced=string_replaced+"Wildcard:"+chunk+"-->"+text+"," 97 | new_prompt=new_prompt+text 98 | else: 99 | new_prompt=new_prompt+text 100 | 101 | return new_prompt 102 | #return new_prompt, string_replaced -------------------------------------------------------------------------------- /UI/config_ui_general.py: -------------------------------------------------------------------------------- 1 | import gradio as gr 2 | import os 3 | 4 | from Engine.General_parameters import UI_Configuration as UI_Configuration 5 | 6 | 7 | models_dir=f"{os.getcwd()}"+'\\models' 8 | 9 | def show_general_configuration(): 10 | ui_config=UI_Configuration() 11 | with gr.Blocks(title="General Config") as ConfigUI: 12 | with gr.Accordion(label="UI Options",open=False): 13 | gr.Markdown("Directories configuracion") 14 | with gr.Row(): 15 | Models_Dir_Select=gr.Textbox(label="Models Directory",lines=1, value=ui_config.models_dir, visible=True, interactive=True) 16 | Output_Path_Select=gr.Textbox(label="Output Directory",lines=1, value=ui_config.output_path, visible=True, interactive=True) 17 | with gr.Row(): 18 | Txt2img_Tab = gr.Checkbox(label="Txt2img Tab",value=ui_config.Txt2img_Tab, info="De/Activate TAB(Applied in next run)") 19 | InPaint_Tab = gr.Checkbox(label="In-Paint Tab",value=ui_config.InPaint_Tab, info="De/Activate TAB(Applied in next run)") 20 | Img2Img_Tab = gr.Checkbox(label="Img2Img Tab",value=ui_config.Img2Img_Tab, info="De/Activate TAB(Applied in next run)") 21 | InstructP2P_Tab = gr.Checkbox(label="InstructP2P Tab",value=ui_config.InstructP2P_Tab, info="De/Activate TAB(Applied in next run)") 22 | Tools_Tab = gr.Checkbox(label="Image Tools Tab",value=ui_config.Tools_Tab, info="De/Activate TAB(Applied in next run)") 23 | ControlNet_Tab = gr.Checkbox(label="ControlNet Tab",value=ui_config.ControlNet_Tab, info="De/Activate TAB(Applied in next run)") 24 | 25 | with gr.Row(): 26 | Advanced_Config = gr.Checkbox(label="Advanced Config",value=ui_config.Advanced_Config, info="Deactivate Avanced Options(Applied in next run)") 27 | UI_NetworkPort=gr.Textbox(label="Gradio Server Port",lines=1, value=ui_config.GradioPort, visible=True, interactive=True) 28 | with gr.Row(): 29 | apply_btn = gr.Button("Apply & Save config", variant="primary") 30 | #loadconfig_btn = gr.Button("Load saved config") 31 | with gr.Row(): 32 | from UI import edit_styles_ui 33 | edit_styles_ui.show_edit_styles_ui() 34 | 35 | UI_options=[Models_Dir_Select,Output_Path_Select,Txt2img_Tab, InPaint_Tab, Img2Img_Tab,InstructP2P_Tab,Tools_Tab, ControlNet_Tab, Advanced_Config,UI_NetworkPort] 36 | apply_btn.click(fn=applyandsave, inputs=UI_options,outputs=None) 37 | #loadconfig_btn.click(fn=loadconfig, inputs=Img2Img_Tab,outputs=Img2Img_Tab) 38 | 39 | def loadconfig(grImg2Img_Tab): 40 | ui_config=UI_Configuration() 41 | ui_config.Txt2img_Tab=False 42 | return gr.Checkbox.update(value=ui_config.Txt2img_Tab) 43 | 44 | 45 | 46 | def applyandsave(Models_Dir_Select,Output_Path_Select,Txt2img_Tab, 47 | InPaint_Tab, Img2Img_Tab,InstructP2P_Tab,Tools_Tab, 48 | ControlNet_Tab, Advanced_Config,UI_NetworkPort): 49 | 50 | ui_config=UI_Configuration() 51 | ui_config.models_dir=Models_Dir_Select 52 | ui_config.output_path=Output_Path_Select 53 | ui_config.Txt2img_Tab=Txt2img_Tab 54 | ui_config.InPaint_Tab=InPaint_Tab 55 | ui_config.Img2Img_Tab=Img2Img_Tab 56 | ui_config.ControlNet_Tab=ControlNet_Tab 57 | ui_config.Tools_Tab=Tools_Tab 58 | ui_config.Advanced_Config=Advanced_Config 59 | ui_config.InstructP2P_Tab = int(InstructP2P_Tab) 60 | ui_config.GradioPort = UI_NetworkPort 61 | 62 | #print(ui_config) 63 | print("Applied and saved, to work with these settings: clean memory and run any pipeline") 64 | ui_config.save_config_json() 65 | 66 | 67 | def Generic_Select_Option(Radio_Select): 68 | config=UI_Configuration() 69 | if Radio_Select == "Yes": 70 | config.wildcards_activated=True 71 | print(params.UI_Configuration().wildcards_activated) 72 | else: 73 | config.wildcards_activated=False 74 | print(params.UI_Configuration().wildcards_activated) 75 | 76 | def load_values(): 77 | return 78 | 79 | -------------------------------------------------------------------------------- /Engine/Pipelines/inpaint_pipeline.py: -------------------------------------------------------------------------------- 1 | from Engine.General_parameters import Engine_Configuration 2 | 3 | 4 | import gc,json 5 | import numpy as np 6 | 7 | from Engine.General_parameters import Engine_Configuration 8 | from Engine import SchedulersConfig 9 | #from Engine.pipelines_engines import SchedulersConfig 10 | from Engine.engine_common_funcs import seed_generator 11 | from Engine import Vae_and_Text_Encoders 12 | 13 | 14 | #optimum pipes 15 | from optimum.onnxruntime import ORTStableDiffusionInpaintPipeline 16 | 17 | class Borg_Inpaint: 18 | _shared_state = {} 19 | def __init__(self): 20 | self.__dict__ = self._shared_state 21 | 22 | 23 | class inpaint_pipe(Borg_Inpaint): 24 | inpaint_pipe = None 25 | seeds = [] 26 | def __init__(self): 27 | Borg_Inpaint.__init__(self) 28 | 29 | def __str__(self): return json.dumps(self.__dict__) 30 | 31 | def initialize(self,model_path,sched_name): 32 | if Vae_and_Text_Encoders().text_encoder == None: 33 | Vae_and_Text_Encoders().load_textencoder(model_path) 34 | if Vae_and_Text_Encoders().vae_decoder == None: 35 | Vae_and_Text_Encoders().load_vaedecoder(model_path) 36 | if Vae_and_Text_Encoders().vae_encoder == None: 37 | Vae_and_Text_Encoders().load_vaeencoder(model_path) 38 | 39 | 40 | provider=Engine_Configuration().MAINPipe_provider['provider'] 41 | provider_options=Engine_Configuration().MAINPipe_provider['provider_options'] 42 | 43 | import onnxruntime as ort 44 | sess_options = ort.SessionOptions() 45 | sess_options.log_severity_level=3 46 | 47 | if self.inpaint_pipe == None: 48 | from optimum.onnxruntime.modeling_ort import ORTModel 49 | from Engine.engine_common_funcs import load_tokenizer_and_config 50 | 51 | print(f"Loading Inpaint Unet session in [{provider}] with options:{provider_options}") 52 | unet_session = ORTModel.load_model(model_path+"/unet/model.onnx", provider,sess_options, provider_options=provider_options) 53 | tokenizer,config=load_tokenizer_and_config(model_path) 54 | 55 | self.inpaint_pipe = ORTStableDiffusionInpaintPipeline( 56 | unet_session=unet_session, 57 | vae_decoder_session= Vae_and_Text_Encoders().vae_decoder, 58 | text_encoder_session= Vae_and_Text_Encoders().text_encoder, 59 | vae_encoder_session=Vae_and_Text_Encoders().vae_encoder, 60 | tokenizer=tokenizer, 61 | config=config, 62 | scheduler=SchedulersConfig().scheduler(sched_name,model_path) 63 | ) 64 | 65 | else: 66 | self.inpaint_pipe.scheduler=SchedulersConfig().scheduler(sched_name,model_path) 67 | return self.inpaint_pipe 68 | 69 | def create_seeds(self,seed=None,iter=1,same_seeds=False): 70 | self.seeds=seed_generator(seed,iter) 71 | if same_seeds: 72 | for seed in self.seeds: 73 | seed = self.seeds[0] 74 | 75 | def unload_from_memory(self): 76 | self.inpaint_pipe= None 77 | self.seeds= None 78 | gc.collect() 79 | 80 | 81 | def run_inference(self,prompt,neg_prompt,init_image,init_mask,height,width,steps,guid,eta,batch,seed): 82 | import numpy as np 83 | rng = np.random.RandomState(seed) 84 | prompt.strip("\n") 85 | neg_prompt.strip("\n") 86 | 87 | 88 | batch_images = self.inpaint_pipe( 89 | prompt, 90 | negative_prompt=neg_prompt, 91 | image=init_image, 92 | mask_image=init_mask, 93 | height=height, 94 | width=width, 95 | num_inference_steps=steps, 96 | guidance_scale=guid, 97 | eta=eta, 98 | num_images_per_prompt=batch, 99 | generator=rng, 100 | ).images 101 | 102 | dictio={'prompt':prompt,'neg_prompt':neg_prompt,'height':height,'width':width,'steps':steps,'guid':guid,'eta':eta,'batch':batch,'seed':seed} 103 | return batch_images,dictio 104 | 105 | -------------------------------------------------------------------------------- /UI/edit_styles_ui.py: -------------------------------------------------------------------------------- 1 | import gradio as gr 2 | from UI import styles_ui 3 | global styles_dict 4 | 5 | def show_edit_styles_ui(): 6 | global styles_dict 7 | styles_dict= get_styles() 8 | styles_keys= list(styles_dict.keys()) 9 | if True: 10 | with gr.Accordion(label="Styles",open=False): 11 | gr.Markdown("Edit your preferred Styles") 12 | with gr.Row(): 13 | with gr.Column(scale=1): 14 | Style_Select = gr.Radio(styles_keys,value=styles_keys[0],label="Available Styles") 15 | with gr.Column(scale=8): 16 | styletext_name = gr.Textbox(value="", lines=1, label="Style name") 17 | styletext_pre = gr.Textbox(value="", lines=2, label="Style previous text") 18 | styletext_post = gr.Textbox(value="", lines=2, label="Style posterior text") 19 | with gr.Row(): 20 | save_btn = gr.Button("Save this Style") 21 | new_style_btn = gr.Button("Create New Blank Style") 22 | del_style_btn = gr.Button("Delete Selected Style") 23 | reload_btn = gr.Button("Reload Styles") 24 | 25 | all_inputs=[Style_Select,styletext_pre,styletext_post,styletext_name] 26 | del_style_btn.click(fn=delete_style, inputs=Style_Select, outputs=Style_Select) 27 | save_btn.click(fn=save_styles, inputs=all_inputs, outputs=Style_Select) 28 | Style_Select.change(fn=select_style, inputs=Style_Select, outputs=[styletext_name,styletext_pre,styletext_post,Style_Select]) 29 | new_style_btn.click(fn=add_new_style, inputs=all_inputs, outputs=Style_Select) 30 | reload_btn.click(fn=reload_styles, inputs=None, outputs=Style_Select) 31 | 32 | 33 | 34 | def reload_styles(*args): 35 | global styles_dict 36 | styles_dict=get_styles() 37 | styles_keys= list(styles_dict.keys()) 38 | return gr.Radio.update(choices=styles_keys,value="None") 39 | 40 | 41 | def add_new_style(*args): 42 | global styles_dict 43 | styles_dict.update({'NewStyle':" | "}) 44 | return Update_StyleSelect() 45 | 46 | def get_styles(): 47 | import json 48 | """dict={ 49 | "None":True, 50 | "StudioPhoto":"(RAW, 8k) |, studio lights,pseudo-impasto", 51 | "Style1":"(cartoon) |, Ink drawing line art", 52 | "Style2":"unity wallpaper, 8k, high quality, | masterpiece,(masterpiece, top quality, best quality)" 53 | }""" 54 | with open('./Engine/config_files/Styles.json', 'r') as openfile: 55 | jsonStr = json.load(openfile) 56 | #Aqui añadir por si es la primera ejecucion y no existe el fichero") 57 | 58 | jsonStr.update({"None":True}) 59 | return jsonStr 60 | 61 | 62 | def select_style(*args): 63 | global styles_dict 64 | dict=styles_dict 65 | style=args[0] 66 | 67 | if style != "None": 68 | params=dict[style].split("|") 69 | return style,params[0],params[1], Update_StyleSelect() 70 | else: 71 | return "None","","",gr.Radio.update(visible=True) 72 | 73 | def delete_style(*args): 74 | import json 75 | global styles_dict 76 | styles_dict 77 | 78 | style=args[0] 79 | 80 | if style != "None": 81 | styles_dict.pop(style) 82 | jsonStr = json.dumps(styles_dict) 83 | with open("./Engine/config_files/Styles.json", "w") as outfile: 84 | outfile.write(jsonStr) 85 | print("Saving Styles without this Style") 86 | else: 87 | print("Cannot Delete the empty Style") 88 | 89 | return Update_StyleSelect() 90 | 91 | def Update_StyleSelect(*args): 92 | global styles_dict 93 | styles_keys= list(styles_dict.keys()) 94 | return gr.Radio.update(choices=styles_keys) 95 | 96 | def save_styles(*args): 97 | import json 98 | global styles_dict 99 | styles_dict 100 | 101 | style=args[0] 102 | style_pre=args[1] 103 | style_post=args[2] 104 | style_name=args[3] 105 | 106 | if style != "None": 107 | styles_dict.pop(style) 108 | styles_dict.update({style_name:f"{style_pre}|{style_post}"}) 109 | jsonStr = json.dumps(styles_dict) 110 | with open("./Engine/config_files/Styles.json", "w") as outfile: 111 | outfile.write(jsonStr) 112 | print("Saving Style") 113 | 114 | return Update_StyleSelect() 115 | 116 | 117 | -------------------------------------------------------------------------------- /UI/tools_ui/clipskip_config_ui.py: -------------------------------------------------------------------------------- 1 | import gradio as gr 2 | from UI import styles_ui 3 | global styles_dict 4 | 5 | def show_edit_styles_ui(): 6 | global styles_dict 7 | styles_dict= get_styles() 8 | styles_keys= list(styles_dict.keys()) 9 | if True: 10 | with gr.Accordion(label="Styles",open=False): 11 | gr.Markdown("Edit your preferred Styles") 12 | with gr.Row(): 13 | with gr.Column(scale=1): 14 | Style_Select = gr.Radio(styles_keys,value=styles_keys[0],label="Available Styles") 15 | with gr.Column(scale=8): 16 | styletext_name = gr.Textbox(value="", lines=1, label="Style name") 17 | styletext_pre = gr.Textbox(value="", lines=2, label="Style previous text") 18 | styletext_post = gr.Textbox(value="", lines=2, label="Style posterior text") 19 | with gr.Row(): 20 | save_btn = gr.Button("Save this Style") 21 | new_style_btn = gr.Button("Create New Blank Style") 22 | del_style_btn = gr.Button("Delete Selected Style") 23 | reload_btn = gr.Button("Reload Styles") 24 | 25 | all_inputs=[Style_Select,styletext_pre,styletext_post,styletext_name] 26 | del_style_btn.click(fn=delete_style, inputs=Style_Select, outputs=Style_Select) 27 | save_btn.click(fn=save_styles, inputs=all_inputs, outputs=Style_Select) 28 | Style_Select.change(fn=select_style, inputs=Style_Select, outputs=[styletext_name,styletext_pre,styletext_post,Style_Select]) 29 | new_style_btn.click(fn=add_new_style, inputs=all_inputs, outputs=Style_Select) 30 | reload_btn.click(fn=reload_styles, inputs=None, outputs=Style_Select) 31 | 32 | 33 | 34 | def reload_styles(*args): 35 | global styles_dict 36 | styles_dict=get_styles() 37 | styles_keys= list(styles_dict.keys()) 38 | return gr.Radio.update(choices=styles_keys,value="None") 39 | 40 | 41 | def add_new_style(*args): 42 | global styles_dict 43 | styles_dict.update({'NewStyle':" | "}) 44 | return Update_StyleSelect() 45 | 46 | def get_styles(): 47 | import json 48 | """dict={ 49 | "None":True, 50 | "StudioPhoto":"(RAW, 8k) |, studio lights,pseudo-impasto", 51 | "Style1":"(cartoon) |, Ink drawing line art", 52 | "Style2":"unity wallpaper, 8k, high quality, | masterpiece,(masterpiece, top quality, best quality)" 53 | }""" 54 | with open('./Engine/config_files/Styles.json', 'r') as openfile: 55 | jsonStr = json.load(openfile) 56 | #Aqui añadir por si es la primera ejecucion y no existe el fichero") 57 | 58 | jsonStr.update({"None":True}) 59 | return jsonStr 60 | 61 | 62 | def select_style(*args): 63 | global styles_dict 64 | dict=styles_dict 65 | style=args[0] 66 | 67 | if style != "None": 68 | params=dict[style].split("|") 69 | return style,params[0],params[1], Update_StyleSelect() 70 | else: 71 | return "None","","",gr.Radio.update(visible=True) 72 | 73 | def delete_style(*args): 74 | import json 75 | global styles_dict 76 | styles_dict 77 | 78 | style=args[0] 79 | 80 | if style != "None": 81 | styles_dict.pop(style) 82 | jsonStr = json.dumps(styles_dict) 83 | with open("./Engine/config_files/Styles.json", "w") as outfile: 84 | outfile.write(jsonStr) 85 | print("Saving Styles without this Style") 86 | else: 87 | print("Cannot Delete the empty Style") 88 | 89 | return Update_StyleSelect() 90 | 91 | def Update_StyleSelect(*args): 92 | global styles_dict 93 | styles_keys= list(styles_dict.keys()) 94 | return gr.Radio.update(choices=styles_keys) 95 | 96 | def save_styles(*args): 97 | import json 98 | global styles_dict 99 | styles_dict 100 | 101 | style=args[0] 102 | style_pre=args[1] 103 | style_post=args[2] 104 | style_name=args[3] 105 | 106 | if style != "None": 107 | styles_dict.pop(style) 108 | styles_dict.update({style_name:f"{style_pre}|{style_post}"}) 109 | jsonStr = json.dumps(styles_dict) 110 | with open("./Engine/config_files/Styles.json", "w") as outfile: 111 | outfile.write(jsonStr) 112 | print("Saving Style") 113 | 114 | return Update_StyleSelect() 115 | 116 | 117 | -------------------------------------------------------------------------------- /modules/styles_module.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | #### MODULE NECESSARY FUNCTIONS (__init__ , show and __call__) #### 4 | def __init__(*args): 5 | __name__='StylesModule' 6 | print(args[0]) 7 | #here a check of access of initializacion if needed 8 | #return (["txt2img","hires"],"prompt_process") 9 | # must return a dict: tabs: in which tabs is to be shown,ui_position: area within the UI tab and func_processing where is to be processed the data 10 | return { 11 | 'tabs':["txt2img","hires"], 12 | 'ui_position':"prompt_process", 13 | 'func_processing':'prompt_process'} 14 | 15 | def is_global_ui(): 16 | return False 17 | 18 | def is_global_function(): 19 | return False 20 | 21 | def show(): 22 | show_styles_ui() 23 | 24 | def __call__(datos): 25 | #what to do when the module is call __call__ 26 | #print("Processing Styles Module2") 27 | if type(datos)==dict: 28 | #print("Entrando Styles como dict") 29 | prompt = datos['prompt_t0'] 30 | prompt = process_prompt(prompt) 31 | datos.update({'prompt_t0':prompt}) 32 | elif type(datos)==str: 33 | #print("Entrando Styles como str") 34 | datos=process_prompt(datos) 35 | else: 36 | print("Not recognized input for module Styles %s" % type(datos)) 37 | datos=None 38 | 39 | return datos 40 | 41 | 42 | 43 | ##### MAIN MODULE CODE ##### 44 | 45 | import gradio as gr 46 | global styles_dict,style 47 | 48 | def show_styles_ui(): 49 | global styles_dict 50 | global style 51 | style=None 52 | styles_dict= get_styles() 53 | styles_keys= list(styles_dict.keys()) 54 | if True: 55 | with gr.Accordion(label="Styles",open=False): 56 | gr.Markdown("Use your preferred Styles") 57 | with gr.Row(): 58 | Style_Select = gr.Radio(styles_keys,value=styles_keys[0],label="Available Styles") 59 | with gr.Row(): 60 | with gr.Accordion(label="Style - modificate running config",open=False): 61 | styletext_pre = gr.Textbox(value="", lines=2, label="Style previous text") 62 | styletext_post = gr.Textbox(value="", lines=2, label="Style posterior text") 63 | with gr.Row(): 64 | apply_btn = gr.Button("Apply Modified Style") 65 | reload_btn = gr.Button("Reload Styles") 66 | 67 | all_inputs=[Style_Select,styletext_pre,styletext_post] 68 | 69 | apply_btn.click(fn=saveto_memory_style, inputs=all_inputs, outputs=None) 70 | reload_btn.click(fn=reload_styles, inputs=None, outputs=Style_Select) 71 | Style_Select.change(fn=apply_styles, inputs=Style_Select, outputs=[styletext_pre,styletext_post]) 72 | 73 | def get_styles(): 74 | import json 75 | with open('./Engine/config_files/Styles.json', 'r') as openfile: 76 | jsonStr = json.load(openfile) 77 | 78 | jsonStr.update({"None":True}) 79 | return jsonStr 80 | 81 | def reload_styles(*args): 82 | global styles_dict 83 | styles_dict=get_styles() 84 | styles_keys= list(styles_dict.keys()) 85 | apply_styles("None") 86 | return gr.Radio.update(choices=styles_keys,value="None") 87 | 88 | def apply_styles(*args): 89 | global styles_dict,style 90 | dict=styles_dict 91 | style_selected=args[0] 92 | style=False 93 | 94 | if style_selected != "None": 95 | style=dict[style_selected] 96 | params=dict[style_selected].split("|") 97 | return params[0],params[1] 98 | else: 99 | return "","" 100 | 101 | def saveto_memory_style(*args): 102 | global styles_dict,style 103 | styles_dict 104 | 105 | style_selected=args[0] 106 | style_pre=args[1] 107 | style_post=args[2] 108 | 109 | style=False 110 | 111 | if style_selected != "None": 112 | styles_dict.update({style_selected:f"{style_pre}|{style_post}"}) 113 | style=styles_dict[style_selected] 114 | 115 | def process_prompt(prompt): 116 | style_pre ="" 117 | style_post="" 118 | global style 119 | 120 | style2=style 121 | if style2: 122 | styles=style2.split("|") 123 | style_pre =styles[0] 124 | style_post=styles[1] 125 | 126 | return style_pre+" " +prompt+" " +style_post 127 | 128 | if __name__ == "__main__": 129 | print("This is a module not intended to run as standalone") 130 | pass 131 | -------------------------------------------------------------------------------- /UI/config_ui_ControlNet.py: -------------------------------------------------------------------------------- 1 | import gradio as gr 2 | from Engine.General_parameters import ControlNet_config 3 | 4 | def show_controlnet_models_configuration(): 5 | ControlNetConfig=ControlNet_config().config 6 | 7 | if True: 8 | with gr.Accordion(label="Select ControlNet Models Directories",open=False): 9 | gr.Markdown("Instead of saving duplicates of each ControlNet model for every model, save one instance of generic ControlNet models into a directory and write down here their path.Include full name and extension") 10 | with gr.Row(): 11 | with gr.Column(scale=1): 12 | canny_active = gr.Checkbox(label="Canny Model Activated?", value=ControlNetConfig["canny_active"], interactive=True) 13 | with gr.Column(scale=8): 14 | canny_path=gr.Textbox(label="Canny model full path",lines=1, value=ControlNetConfig["canny_path"], visible=True, interactive=True) 15 | with gr.Row(): 16 | with gr.Column(scale=1): 17 | depth_active = gr.Checkbox(label="Depth Model Activated?", value=ControlNetConfig["depth_active"], interactive=True) 18 | with gr.Column(scale=8): 19 | depth_path=gr.Textbox(label="Depth Model full path",lines=1, value=ControlNetConfig["depth_path"], visible=True, interactive=True) 20 | with gr.Row(): 21 | with gr.Column(scale=1): 22 | hed_active = gr.Checkbox(label="Hed Model Activated?", value=ControlNetConfig["hed_active"], interactive=True) 23 | with gr.Column(scale=8): 24 | hed_path=gr.Textbox(label="Hed Model full path",lines=1, value=ControlNetConfig["hed_path"], visible=True, interactive=True) 25 | with gr.Row(): 26 | with gr.Column(scale=1): 27 | mlsd_active = gr.Checkbox(label="Mlsd Model Activated?", value=ControlNetConfig["mlsd_active"], interactive=True) 28 | with gr.Column(scale=8): 29 | mlsd_path=gr.Textbox(label="Mlsd Model full path",lines=1, value=ControlNetConfig["mlsd_path"], visible=True, interactive=True) 30 | with gr.Row(): 31 | with gr.Column(scale=1): 32 | normal_active = gr.Checkbox(label="Normal Model Activated?", value=ControlNetConfig["normal_active"], interactive=True) 33 | with gr.Column(scale=8): 34 | normal_path=gr.Textbox(label="Normal Model full path",lines=1, value=ControlNetConfig["normal_path"], visible=True, interactive=True) 35 | with gr.Row(): 36 | with gr.Column(scale=1): 37 | openpose_active = gr.Checkbox(label="Openpose Model Activated?", value=ControlNetConfig["openpose_active"], interactive=True) 38 | with gr.Column(scale=8): 39 | openpose_path=gr.Textbox(label="Openpose full path",lines=1, value=ControlNetConfig["openpose_path"], visible=True, interactive=True) 40 | with gr.Row(): 41 | with gr.Column(scale=1): 42 | seg_active = gr.Checkbox(label="Seg Model Activated?", value=ControlNetConfig["seg_active"], interactive=True) 43 | with gr.Column(scale=8): 44 | seg_path=gr.Textbox(label="Seg Model full path",lines=1, value=ControlNetConfig["seg_path"], visible=True, interactive=True) 45 | with gr.Row(): 46 | save_btn = gr.Button("Apply & Save ControlNet models config") 47 | load_btn = gr.Button("Load ControlNet models config") 48 | 49 | all_inputs=[ 50 | canny_active,canny_path,depth_active,depth_path,hed_active,hed_path,mlsd_active,mlsd_path, 51 | normal_active, normal_path,openpose_active,openpose_path,seg_active,seg_path] 52 | 53 | save_btn.click(fn=save_controlnet_config_ui, inputs=all_inputs, outputs=None) 54 | load_btn.click(fn=load_controlnet_config_ui, inputs=None , outputs=all_inputs) 55 | 56 | 57 | def load_controlnet_config_ui(): 58 | config=ControlNet_config() 59 | config.load_config_from_disk() 60 | return list(ControlNet_config().config.values()) 61 | 62 | 63 | def save_controlnet_config_ui(*args): 64 | controlnet_config= { 65 | "canny_active":args[0], 66 | "canny_path":args[1], 67 | "depth_active":args[2], 68 | "depth_path":args[3], 69 | "hed_active":args[4], 70 | "hed_path":args[5], 71 | "mlsd_active":args[6], 72 | "mlsd_path":args[7], 73 | "normal_active":args[8], 74 | "normal_path":args[9], 75 | "openpose_active":args[10], 76 | "openpose_path":args[11], 77 | "seg_active":args[12], 78 | "seg_path":args[13], 79 | } 80 | ControlNet_config().save_controlnet_config(controlnet_config) 81 | 82 | 83 | 84 | -------------------------------------------------------------------------------- /Scripts/utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | 4 | 5 | def get_names(inp, prefix): 6 | if not isinstance(inp, (tuple, list)): 7 | inp = [inp] 8 | 9 | names = [] 10 | for i, sub_inp in enumerate(inp): 11 | sub_prefix = '{}.{}'.format(prefix, i) 12 | if isinstance(sub_inp, (list, tuple)): 13 | names.extend(get_names(sub_inp, sub_prefix)) 14 | else: 15 | names.append(sub_prefix) 16 | 17 | return names 18 | 19 | 20 | def get_forms(inp): 21 | if not isinstance(inp, (tuple, list)): 22 | return 'x' 23 | 24 | forms = [] 25 | for sub_inp in inp: 26 | forms.append(get_forms(sub_inp)) 27 | 28 | return forms 29 | 30 | 31 | def to(inp, device_or_dtype): 32 | if not isinstance(inp, (tuple, list)): 33 | if type(inp).__module__ == torch.__name__: 34 | if device_or_dtype == 'torch': 35 | pass 36 | elif device_or_dtype == 'numpy': 37 | inp = inp.detach().cpu().numpy() 38 | else: 39 | inp = inp.to(device_or_dtype) 40 | elif type(inp).__module__ == np.__name__: 41 | if not isinstance(inp, np.ndarray): 42 | inp = np.array(inp) 43 | 44 | if device_or_dtype == 'torch': 45 | inp = torch.from_numpy(inp) 46 | elif device_or_dtype == 'numpy': 47 | pass 48 | else: 49 | inp = inp.astype(device_or_dtype) 50 | elif isinstance(inp, (int, float)): 51 | if device_or_dtype == 'torch': 52 | inp = torch.tensor(inp) 53 | elif device_or_dtype == 'numpy': 54 | inp = np.array(inp) 55 | else: 56 | raise TypeError(('Unsupported type {}, expect int, float, ' 57 | 'np.ndarray or torch.Tensor').format(type(inp))) 58 | 59 | return inp 60 | 61 | out = [] 62 | for sub_inp in inp: 63 | out.append(to(sub_inp, device_or_dtype)) 64 | 65 | return out 66 | 67 | 68 | def flatten(inp): 69 | if not isinstance(inp, (tuple, list)): 70 | return [inp] 71 | 72 | out = [] 73 | for sub_inp in inp: 74 | out.extend(flatten(sub_inp)) 75 | 76 | return out 77 | 78 | 79 | def reconstruct(inp, forms): 80 | assert len(flatten(inp)) == len(flatten(forms)) 81 | 82 | if not isinstance(forms, (tuple, list)): 83 | if isinstance(inp, (tuple, list)): 84 | assert len(inp) == 1 85 | return inp[0] 86 | else: 87 | return inp 88 | 89 | out = [] 90 | index = 0 91 | for sub_form in forms: 92 | if isinstance(sub_form, (tuple, list)): 93 | sub_form_len = len(flatten(sub_form)) 94 | out.append(reconstruct(inp[:sub_form_len], sub_form)) 95 | index += sub_form_len 96 | else: 97 | out.append(inp[index]) 98 | index += 1 99 | 100 | return out 101 | 102 | 103 | def add_batch_dim(inp): 104 | if not isinstance(inp, (list, tuple)): 105 | return inp[None, ...] 106 | 107 | out = [] 108 | for sub_inp in inp: 109 | out.append(add_batch_dim(sub_inp)) 110 | 111 | return out 112 | 113 | 114 | def cat(x, y, dim=0): 115 | x = list(x) if isinstance(x, (tuple, list)) else x 116 | y = list(y) if isinstance(y, (tuple, list)) else y 117 | 118 | assert type(x) == type(y) 119 | 120 | if isinstance(x, list): 121 | assert len(x) == len(y) 122 | 123 | out = [] 124 | for sub_x, sub_y in zip(x, y): 125 | out.append(cat(sub_x, sub_y, dim)) 126 | return out 127 | elif isinstance(x, torch.Tensor): 128 | return torch.cat([x, y], dim=dim) 129 | elif isinstance(x, np.ndarray): 130 | return np.concatenate([x, y], axis=dim) 131 | else: 132 | raise TypeError(('Unsupported data type {}, ' 133 | 'expect np.ndarray or torch.Tensor').format(type(x))) 134 | 135 | 136 | def gen_ones(shape): 137 | if isinstance(shape[0], int): 138 | return torch.ones(*shape) 139 | 140 | data = [] 141 | for sub_shape in shape: 142 | data.append(gen_ones(sub_shape)) 143 | 144 | return data 145 | 146 | 147 | def fetch_batch(data, start, end): 148 | if isinstance(data, torch.Tensor): 149 | return data[start:end] 150 | 151 | assert not isinstance(data, (tuple, list)), ( 152 | 'Unsupported data type {}, only torch.Tensor, ' 153 | 'tuple or list are supported').format(type(data)) 154 | 155 | batch = [] 156 | for sub_data in data: 157 | batch.append(fetch_batch(sub_data, start, end)) 158 | 159 | return batch -------------------------------------------------------------------------------- /Engine/Pipelines/img2img_pipeline.py: -------------------------------------------------------------------------------- 1 | #from sched import scheduler 2 | #from Engine.pipelines_engines import SchedulersConfig 3 | from Engine.General_parameters import Engine_Configuration 4 | from Engine import Vae_and_Text_Encoders 5 | from Engine import SchedulersConfig 6 | from Engine.engine_common_funcs import seed_generator 7 | 8 | #optimum pipes 9 | from optimum.onnxruntime import ORTStableDiffusionImg2ImgPipeline 10 | 11 | import gc 12 | import json 13 | import numpy as np 14 | 15 | 16 | 17 | 18 | class Borg_img2img: 19 | _shared_state = {} 20 | def __init__(self): 21 | self.__dict__ = self._shared_state 22 | 23 | 24 | class img2img_pipe(Borg_img2img): 25 | img2img_pipe = None 26 | model = None 27 | seeds = [] 28 | def __init__(self): 29 | Borg_img2img.__init__(self) 30 | 31 | def __str__(self): return json.dumps(self.__dict__) 32 | 33 | def initialize(self,model_path,sched_name): 34 | #from Engine.General_parameters import Engine_Configuration as en_config 35 | if Vae_and_Text_Encoders().text_encoder == None: 36 | Vae_and_Text_Encoders().load_textencoder(model_path) 37 | if Vae_and_Text_Encoders().vae_decoder == None: 38 | Vae_and_Text_Encoders().load_vaedecoder(model_path) 39 | if Vae_and_Text_Encoders().vae_encoder == None: 40 | Vae_and_Text_Encoders().load_vaeencoder(model_path) 41 | 42 | provider=Engine_Configuration().MAINPipe_provider['provider'] 43 | provider_options=Engine_Configuration().MAINPipe_provider['provider_options'] 44 | 45 | 46 | import onnxruntime as ort 47 | sess_options = ort.SessionOptions() 48 | sess_options.log_severity_level=3 49 | 50 | unet_session=load_unet_model(model_path, provider=provider, sess_options=sess_options, provider_options=provider_options) 51 | tokenizer,config=load_tokenizer(model_path) 52 | 53 | if self.img2img_pipe == None: 54 | print(f"Loading Img2Img pipe in {provider} with options: {provider_options}") 55 | 56 | self.img2img_pipe =ORTStableDiffusionImg2ImgPipeline( 57 | unet_session=unet_session, 58 | tokenizer=tokenizer, 59 | config=config, 60 | scheduler=SchedulersConfig().scheduler(sched_name,model_path), 61 | text_encoder_session=Vae_and_Text_Encoders().text_encoder, 62 | vae_decoder_session=Vae_and_Text_Encoders().vae_decoder, 63 | vae_encoder_session=Vae_and_Text_Encoders().vae_encoder, 64 | ) 65 | else: 66 | self.img2img_pipe.scheduler=SchedulersConfig().scheduler(sched_name,model_path) 67 | return self.img2img_pipe 68 | 69 | def create_seeds(self,seed=None,iter=1,same_seeds=False): 70 | self.seeds=seed_generator(seed,iter) 71 | if same_seeds: 72 | for seed in self.seeds: 73 | seed = self.seeds[0] 74 | 75 | def unload_from_memory(self): 76 | self.img2img_pipe= None 77 | self.model = None 78 | self.seeds = None 79 | gc.collect() 80 | 81 | 82 | def run_inference(self,prompt,neg_prompt,init_image,strength,steps,guid,eta,batch,seed): 83 | import numpy as np 84 | rng = np.random.RandomState(seed) 85 | prompt.strip("\n") 86 | neg_prompt.strip("\n") 87 | 88 | 89 | batch_images = self.img2img_pipe( 90 | prompt, 91 | negative_prompt=neg_prompt, 92 | image=init_image, 93 | strength= strength, 94 | num_inference_steps=steps, 95 | guidance_scale=guid, 96 | eta=eta, 97 | num_images_per_prompt=batch, 98 | generator=rng, 99 | ).images 100 | dictio={'Img2ImgPrompt':prompt,'neg_prompt':neg_prompt,'steps':steps,'guid':guid,'eta':eta,'strength':strength,'seed':seed} 101 | return batch_images,dictio 102 | 103 | 104 | def load_unet_model(model_path,provider, sess_options, provider_options,model_name: str = None): 105 | from optimum.onnxruntime.modeling_ort import ORTModel 106 | if model_name==None: 107 | model_name="/unet/model.onnx" 108 | return ORTModel.load_model(model_path+model_name, provider,sess_options, provider_options=provider_options) 109 | 110 | 111 | def load_tokenizer(model_path): 112 | import json 113 | CONFIG_NAME="model_index.json" 114 | config_file=model_path+"/"+CONFIG_NAME 115 | config=None 116 | try: 117 | import json 118 | with open(config_file, "r", encoding="utf-8") as reader: 119 | text = reader.read() 120 | config=json.loads(text) 121 | except: 122 | raise OSError(f"model_index.json not found in {model_path} local folder") 123 | 124 | from transformers import CLIPTokenizer 125 | return CLIPTokenizer.from_pretrained(model_path+"/"+'tokenizer'),config -------------------------------------------------------------------------------- /ONNX-StableDiffusion.py: -------------------------------------------------------------------------------- 1 | import gradio as gr 2 | from Engine.General_parameters import Engine_Configuration 3 | from Engine.General_parameters import UI_Configuration 4 | from Engine.General_parameters import running_config 5 | 6 | from UI import UI_placement_areas 7 | from modules.modules import preProcess_modules 8 | global list_modules 9 | list_modules=[] 10 | list_modules=preProcess_modules().check_available_modules("main_ui") 11 | 12 | 13 | def init_ui(): 14 | ui_config=UI_Configuration() 15 | with gr.Blocks(title="ONNX Difussers Modular UI",css= css1) as demo: 16 | if ui_config.Txt2img_Tab: 17 | with gr.Tab(label="Testing HiRes Txt2img Pipeline") as tab10: 18 | from UI import HiRes_txt2img_ui as HiRes_txt2img_ui 19 | HiRes_txt2img_ui.show_HiRes_txt2img_ui() 20 | if ui_config.Txt2img_Tab: 21 | with gr.Tab(label="Txt2img Pipelines & Inferences") as tab0: 22 | from UI import txt2img_ui as txt2img_ui 23 | txt2img_ui.show_txt2img_ui() 24 | if ui_config.Img2Img_Tab: 25 | with gr.Tab(label="Img2Img") as tab1: 26 | from UI import Img2Img_ui 27 | Img2Img_ui.show_Img2Img_ui() 28 | if ui_config.InPaint_Tab: 29 | with gr.Tab(label="InPaint") as tab2: 30 | from UI import Inpaint_ui 31 | Inpaint_ui.show_Inpaint_ui() 32 | #if ui_config.Tools_Tab: 33 | if False: 34 | with gr.Tab(label="Image Tools") as tab3: 35 | from UI import ui_image_tools 36 | ui_image_tools.show_input_image_area() 37 | ui_image_tools.show_danbooru_area() 38 | ui_image_tools.show_image_resolution_area() 39 | #if ui_config.InstructP2P_Tab: 40 | if False: 41 | with gr.Tab(label="Instruct Pix2Pix") as tab4: 42 | from UI import instructp2p_ui 43 | instructp2p_ui.show_instructp2p_ui() 44 | if ui_config.ControlNet_Tab: 45 | with gr.Tab(label="ControlNet") as tab5: 46 | from UI import ControlNet_ui 47 | ControlNet_ui.show_ControlNet_ui() 48 | if True: 49 | with gr.Tab(label="FaceRestoration") as tab6: 50 | from UI import ui_face_tools 51 | ui_face_tools.show_input_image_area() 52 | 53 | with gr.Tab(label="Configuration") as tab7: 54 | from UI import config_ui_general 55 | from UI import config_ui_ControlNet 56 | from UI import config_ui_Vae 57 | from UI import config_ui_TextEncoder 58 | 59 | config_ui_general.show_general_configuration() 60 | if ui_config.Advanced_Config: 61 | from UI import config_ui_engine as config_ui_engine 62 | config_ui_engine.show_providers_configuration() 63 | config_ui_ControlNet.show_controlnet_models_configuration() 64 | config_ui_TextEncoder.show_textenc_models_configuration() 65 | config_ui_Vae.show_vae_models_configuration() 66 | #from UI import config_ui_wildcards as wilcards_ui_config 67 | #wilcards_ui_config.show_wilcards_configuration() 68 | 69 | ######## Show new ta areas ############################# 70 | #UI_placement_areas.show_new_tabs(list_modules,"main_ui") 71 | ######################################################## 72 | 73 | 74 | ######## Show footer areas ############################# 75 | with gr.Row(): 76 | UI_placement_areas.show_footer_area(list_modules,"main_ui") 77 | ######################################################## 78 | return demo 79 | 80 | 81 | css1 = """ 82 | #title1 {background-color: #00ACAA;text-transform: uppercase; font-weight: bold;} 83 | .feedback textarea {font-size: 24px !important} 84 | #imagen1 {min-height: 600px} 85 | #imagen2 {height: 250px} 86 | #imagen3 {min-width: 512px,min-height: 512px} 87 | """ 88 | 89 | css2 = """ 90 | #title1 {background-color: #00ACAA;text-transform: uppercase; font-weight: bold;} 91 | .feedback textarea {font-size: 24px !important} 92 | """ 93 | 94 | 95 | Running_information= running_config().Running_information 96 | Running_information.update({"cancelled":False}) 97 | Running_information.update({"model":""}) 98 | Running_information.update({"ControlNetModel":""}) #check if used 99 | Running_information.update({"tab":""}) 100 | Running_information.update({"Running":False}) 101 | Running_information.update({"Save_Latents":False}) 102 | Running_information.update({"Load_Latents":False}) 103 | Running_information.update({"Latent_Name":""}) 104 | Running_information.update({"Latent_Formula":""}) 105 | Running_information.update({"Callback_Steps":2}) 106 | Running_information.update({"Vae_Config":["model"]*6}) 107 | Running_information.update({"Textenc_Config":["model"]*2}) 108 | Running_information.update({"offset":1}) 109 | Running_information.update({"Style":False}) 110 | 111 | 112 | Engine_Configuration().load_config_json() 113 | demo =init_ui() 114 | demo.launch(server_port=UI_Configuration().GradioPort) 115 | 116 | 117 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # ONNX-ModularUI-StableDiffusion 2 | This is a working ONNX version of a UI for Stable Diffusion using optimum pipelines. Leveraging ONNX runtime environment for faster inference, working on most common GPU vendors: NVIDIA,AMD GPU...as long as they got support into onnxruntime. Works on low profile 4Gb GPU cards ( and also CPU only, but i did not tested its performance) 3 | 4 | 5 | **Now including TXT2IMG, IMG2IMG,ControlNet, Inpaint, HIRES approach of TXT2IMG and Latent Consistency** 6 | 7 | **Updates:** 8 | Jan 2024 9 | Added Latent Consistency pipelines and support for Txt2Img and as Low-Res Model for Hi-Res approach (LCM Scheduler to be selected if using) 10 | Move some tasks to modules approach 11 | Minor Bugs Solving & Details added 12 | Small tool to check latents directory (root:latent_to_pil.py) 13 | 14 | Dec 2023 15 | Added modules support 16 | Added Model Conversion interface (independent file in ConversionTool:ONNX-SD-ModelConverter.py) 17 | Added two models approach for hi-res pipeline. 18 | Added facerestoration through faceswapping 19 | Added metadata strip and mask selection, independent file(root): IMG-Strip.py 20 | 21 | Nov 2023: 22 | Implemented ControlNet support 23 | Implemented img2img support 24 | Implemented inpaint support 25 | Implemented txt2img support 26 | Implemented hi-res txt2img support 27 | 28 | In dev & Coming: 29 | Module: Library for models & saving favorite prompts (with seeds) 30 | ------------------------------------ 31 | 32 | Already supported: wildcards,pre-defined styles, latent import for hires-txt2img and txt2img inferences, HiRes supports using same or different models for low and hires inferences... 33 | Adapted version of my approach for a working version of stabble diffusion for onnx, able to run on low profile GPUs (also in high-profile), dedicated GPU memory needed:4Gb. 34 | 35 | HiRes approach: 36 | You might use one model to create the composition and a different model to refine the style and details: 37 | ![example](https://github.com/NeusZimmer/ONNX-ModularUI-StableDiffusion/assets/94193584/14722f28-811c-4bf0-a7e7-231d8356a072) 38 | 39 | 40 | 41 | # Install (20-30 min, difficulty low/medium): 42 | 43 | At first, download or clone this repository: ```git clone https://github.com/NeusZimmer/ONNX-ModularUI-StableDiffusion``` 44 | Then, as usual, create a python virtual environment for the install ( recommended, if you want to do it over you base python install that's up to you) 45 | Python version: 3.10 ( tested, please feedback if you make it run in 3.11) 46 | ``` 47 | "cd ONNX-ModularUI-StableDiffusion" 48 | "python -m venv venv" 49 | ``` 50 | #(or ```py -3.10 -m venv venv ``` if running multiple python instances on the same machine) 51 | 52 | Update pip: 53 | ```pip install --upgrade pip``` 54 | 55 | Install requirements file: 56 | ```pip install -r requirements.txt``` 57 | 58 | 59 | Optional, for face restoration/swapping verssion ( recommended to be done next to finalizing the installation of the other and you check everything else is already working) 60 | ```pip install -r additional_requirements.txt``` 61 | In some cases, you may need to install VS c++ 14.0 build tool for insightface: https://aka.ms/vs/17/release/vs_BuildTools.exe 62 | 63 | When finished, I recommend to reinstall again the onnxruntime package, only one of the listed below, according to your GPU vendor, sometimes it did not recognize the adecuate ExecutionProviders (i.e.DMLExecutionProvider) until this package is reinstalled" 64 | (Select this last package install according to onnx documentation for NVIDIA, Intel Vino, Mac...more info about what package you need to install: https://onnxruntime.ai/docs/install/#python-installs) 65 | 66 | ``` 67 | #Attention: ONLY ONE OF THE LIST BELOW: 68 | #AMD (and other DirectMl supported cards): 69 | pip install onnxruntime-directml 70 | #CPU: 71 | pip install onnxruntime 72 | #CUDA-TENSORRT 73 | pip install onnxruntime-gpu 74 | #OpenVino 75 | pip install onnxruntime-openvino 76 | pip install openvino 77 | #Jetson 78 | #Here: https://elinux.org/Jetson_Zoo#ONNX_Runtim 79 | #Azure: 80 | pip install onnxruntime-azure 81 | ``` 82 | Optional: download the onnx model of your preference and include it into the models folder ( or configure the directory of your preference in the Configuration Tab). 83 | **Sample Models** 84 | ``` 85 | https://civitai.com/models/125580/onnx-base-set-of-model-neusz 86 | https://huggingface.co/Neus/CommonModels 87 | ``` 88 | (You will need to download the main model, the vae decoder, vae encoder and text encoder and put each one in the adequate subdir) 89 | 90 | Activate the environment ```activate.bat``` 91 | and the run the UI: ```py -O ONNX-StableDiffusion.py``` 92 | or ```run.bat``` 93 | 94 | Point your browser to localhost:7860 to accessing the UI app and enjoy yourself... 95 | 96 | # Utility for Model Conversion: 97 | (conversion code extracted from: https://github.com/Amblyopius/Stable-Diffusion-ONNX-FP16) 98 | ``` 99 | cd ConversionTool 100 | py ONNX-SD-ModelConverter.py 101 | ``` 102 | And point your browser to localhost:7860 103 | 104 | 105 | Point to old version ( works up to diffusers 14.0):https://github.com/NeusZimmer/ONNX-Stable-Diffusion-ModularUI 106 | (TXT2IMG, HIRES-TXT2IMG, IMG2IMG, ControlNet, InstructPix, 3 approaches for face/image restoration, danbooru tagging ...and more) 107 | 108 | 109 | 110 | 111 | 112 | -------------------------------------------------------------------------------- /modules/image_to_numpy_module.py: -------------------------------------------------------------------------------- 1 | 2 | import gradio as gr 3 | 4 | 5 | #### MODULE NECESSARY FUNCTIONS (__init__ , show and __call__) #### 6 | def __init__(*args): 7 | __name__='ImageToNumpyModule' 8 | print(args[0]) 9 | #here a check of access of initializacion if needed 10 | # must return a dict: tabs: in which tabs is to be shown,ui_position: area within the UI tab and func_processing where is to be processed the data 11 | return { 12 | 'tabs':["hires"], 13 | 'ui_position':"tools_area", 14 | 'func_processing':None} 15 | 16 | def is_global_ui(): 17 | return False 18 | 19 | def is_global_function(): 20 | return False 21 | 22 | def show(): 23 | show_ImageToNumpyModule_ui() 24 | 25 | def __call__(datos): 26 | return datos 27 | 28 | 29 | 30 | ##### MAIN MODULE CODE ##### 31 | 32 | 33 | def show_ImageToNumpyModule_ui(*args): 34 | args= [""] if not args else args 35 | 36 | with gr.Accordion(label="Process IMG To latent",open=False): 37 | with gr.Row(): 38 | image_in = gr.Image(label="input image", type="pil", elem_id="image_init") 39 | with gr.Row(): 40 | height = gr.Slider(64, 2048, value=512, step=64, label="Output Height") 41 | width = gr.Slider(64, 2048, value=512, step=64, label="Output Width") 42 | with gr.Row(): 43 | convert_to_latent_btn = gr.Button("Convert IMG to latent/numpy on disk") 44 | analyze_btn = gr.Button("Show info on image and conversion") 45 | with gr.Row(): 46 | information = gr.Textbox(value=args[0], lines=1, label="Information",interactive=False) 47 | 48 | 49 | convert_to_latent_btn.click(fn=convert_click,inputs=[image_in,height,width],outputs=information) 50 | analyze_btn.click(fn=get_info,inputs=[image_in,height,width],outputs=information) 51 | 52 | return 53 | 54 | def get_info(image,height,width): 55 | image_size=image.size 56 | image_ratio= str(image_size[0]/image_size[1])[0:4] 57 | 58 | destination_ratio=str(width/height)[0:4] 59 | 60 | return f"Image Size(H/W):{(image_size[1],image_size[0])} is ratio:{image_ratio}, destination ratio:{destination_ratio}" 61 | 62 | def convert_click(image,height,width): 63 | import numpy as np 64 | from Engine.General_parameters import running_config 65 | 66 | Running_information= running_config().Running_information 67 | vaeencoder=False 68 | 69 | if Running_information["tab"] == "hires_txt2img": 70 | from Engine import txt2img_hires_pipe 71 | if txt2img_hires_pipe().hires_pipe!=None: 72 | vaeencoder=txt2img_hires_pipe().hires_pipe.vae_encoder 73 | """elif Running_information["tab"] == "txt2img": #txt2img no tiene encoder usar inpaint o load uno para el proceso? 74 | from Engine import txt2img_pipe 75 | if txt2img_pipe().txt2img_pipe!=None: 76 | vaeencoder=txt2img_pipe().txt2img_pipe.vae_encoder""" 77 | 78 | 79 | image=resize_and_crop(image,height,width) 80 | image=np.array(image) 81 | image = np.array(image).astype(np.float32) / 255.0 82 | image = 2.0 * image - 1.0 83 | image = np.expand_dims(image, axis=0) 84 | image = image.transpose(0, 3, 1, 2) 85 | 86 | # encode the init image into latents and scale the latents 87 | 88 | if vaeencoder==False: 89 | from Engine import Vae_and_Text_Encoders 90 | vaeencoder_session=Vae_and_Text_Encoders().load_vaeencoder("") 91 | try: 92 | init_latents = vaeencoder_session.run(['latent_sample'],{'sample': image})[0] 93 | except: 94 | return ("Initialize a model to process an image or configure a standard vae encoder") 95 | 96 | """ 97 | #Esto es para sacar los nombres para poder ejecutar la sesion. 98 | input_names = {input_key.name: idx for idx, input_key in enumerate(vaeencoder_session.get_inputs())} 99 | print(input_names) 100 | output_names={output_key.name: idx for idx, output_key in enumerate(vaeencoder_session.get_outputs())} 101 | print(output_names) 102 | init_latents = None""" 103 | 104 | else: 105 | init_latents = vaeencoder(sample=image)[0] 106 | 107 | init_latents = 0.18215 * init_latents 108 | path="./latents" 109 | savename=f"{path}/LastGenerated_latent.npy" 110 | np.save(savename, init_latents) 111 | 112 | return f"Converted in:{savename} with size of (H/W):{image.shape[0]}x{image.shape[1]}" 113 | 114 | def resize_and_crop(input_image, height, width): 115 | from PIL import Image 116 | input_width, input_height = input_image.size 117 | 118 | # nearest neighbor for upscaling 119 | if (input_width * input_height) < (width * height): 120 | resample_type = Image.NEAREST 121 | # lanczos for downscaling 122 | else: 123 | resample_type = Image.LANCZOS 124 | 125 | if height / width > input_height / input_width: 126 | adjust_width = int(input_width * height / input_height) 127 | input_image = input_image.resize((adjust_width, height), 128 | resample=resample_type) 129 | left = (adjust_width - width) // 2 130 | right = left + width 131 | input_image = input_image.crop((left, 0, right, height)) 132 | else: 133 | adjust_height = int(input_height * width / input_width) 134 | input_image = input_image.resize((width, adjust_height), 135 | resample=resample_type) 136 | top = (adjust_height - height) // 2 137 | bottom = top + height 138 | input_image = input_image.crop((0, top, width, bottom)) 139 | return input_image 140 | 141 | if __name__ == "__main__": 142 | print("This is a module not intended to run as standalone") 143 | with gr.Blocks(title="Image Converter") as demo: 144 | show_ImageToNumpyModule_ui("This is a module not intended to run as standalone, will fail loading the model for converting the image or saving to disk") 145 | demo.launch() 146 | 147 | pass 148 | -------------------------------------------------------------------------------- /Engine/Pipelines/Controlnet_pipeline.py: -------------------------------------------------------------------------------- 1 | import gc,json 2 | 3 | 4 | from Engine.General_parameters import ( 5 | Engine_Configuration, 6 | ControlNet_config 7 | ) 8 | 9 | from Engine.engine_common_funcs import load_tokenizer_and_config,seed_generator 10 | from Engine import Vae_and_Text_Encoders 11 | from Engine import SchedulersConfig 12 | #from Engine.pipelines_engines import SchedulersConfig 13 | 14 | 15 | from Engine.pipeline_onnx_stable_diffusion_controlnet import OnnxStableDiffusionControlNetPipeline 16 | 17 | 18 | from diffusers.utils.torch_utils import randn_tensor 19 | from diffusers import ( 20 | OnnxRuntimeModel, 21 | OnnxStableDiffusionPipeline, 22 | OnnxStableDiffusionInpaintPipeline, 23 | OnnxStableDiffusionInpaintPipelineLegacy, 24 | OnnxStableDiffusionImg2ImgPipeline, 25 | ) 26 | 27 | class Borg6: 28 | _shared_state = {} 29 | def __init__(self): 30 | self.__dict__ = self._shared_state 31 | 32 | 33 | class ControlNet_pipe(Borg6): 34 | #import onnxruntime as ort 35 | controlnet_Model_ort= None 36 | #controlnet_unet_ort= None 37 | ControlNET_Name=None 38 | ControlNet_pipe = None 39 | unet_model =None 40 | seeds = [] 41 | 42 | def __init__(self): 43 | Borg6.__init__(self) 44 | 45 | def __str__(self): return json.dumps(self.__dict__) 46 | 47 | def load_ControlNet_model(self,model_path,ControlNET_drop): 48 | #self.__load_ControlNet_model(model_path,ControlNET_drop) 49 | self.__load_ControlNet_model(ControlNET_drop) 50 | 51 | def __load_ControlNet_model(self,ControlNET_drop): 52 | import onnxruntime as ort 53 | 54 | provider=Engine_Configuration().ControlNet_provider['provider'] 55 | provider_options=Engine_Configuration().ControlNet_provider['provider_options'] 56 | print(f"Loading ControlNet module at {provider} with options:{provider_options}") 57 | available_models=dict(ControlNet_config().available_controlnet_models()) 58 | 59 | 60 | sess_options = ort.SessionOptions() 61 | sess_options.enable_cpu_mem_arena = False 62 | sess_options.enable_mem_pattern = False 63 | sess_options.log_severity_level=3 64 | self.controlnet_Model_ort= None 65 | 66 | self.ControlNET_Name=ControlNET_drop 67 | ControlNet_path = available_models[ControlNET_drop] 68 | provider=(str(provider),dict(provider_options)) 69 | #self.controlnet_Model_ort = OnnxRuntimeModel.from_pretrained(ControlNet_path, sess_options=opts, provider=provider) 70 | self.controlnet_Model_ort = OnnxRuntimeModel.from_pretrained(ControlNet_path, provider=provider, sess_options=sess_options) 71 | 72 | return self.controlnet_Model_ort 73 | 74 | def __load_uNet_model(self,model_path): 75 | #Aqui cargar con ort el modelo unicamente en el provider principal. 76 | print("Loading Unet module") 77 | 78 | provider=Engine_Configuration().MAINPipe_provider['provider'] 79 | provider_options=Engine_Configuration().MAINPipe_provider['provider_options'] 80 | 81 | import onnxruntime as ort 82 | sess_options = ort.SessionOptions() 83 | sess_options.enable_cpu_mem_arena = False 84 | sess_options.enable_mem_pattern = False 85 | sess_options.log_severity_level=3 86 | provider=(str(provider),dict(provider_options)) 87 | unet_model = OnnxRuntimeModel.from_pretrained(model_path + "/unet", provider=provider, sess_options=sess_options) 88 | return unet_model 89 | 90 | def initialize(self,model_path,sched_name,ControlNET_drop): 91 | if Vae_and_Text_Encoders().text_encoder == None: 92 | Vae_and_Text_Encoders().load_textencoder(model_path,old_version=True) 93 | if Vae_and_Text_Encoders().vae_decoder == None: 94 | Vae_and_Text_Encoders().load_vaedecoder(model_path,old_version=True) 95 | if Vae_and_Text_Encoders().vae_encoder == None: 96 | Vae_and_Text_Encoders().load_vaeencoder(model_path,old_version=True) 97 | 98 | if self.unet_model == None: 99 | self.unet_model=self.__load_uNet_model(model_path) 100 | 101 | #print("Loading Tokenizer") 102 | self.tokenizer,self.config=load_tokenizer_and_config(model_path) 103 | 104 | 105 | if self.ControlNet_pipe == None: 106 | print(f"Using modified model for ControlNET:{model_path}") 107 | self.controlnet_Model_ort= self.__load_ControlNet_model(ControlNET_drop) 108 | 109 | self.ControlNet_pipe = OnnxStableDiffusionControlNetPipeline( #cambiar desde pretrained a __init__ 110 | vae_encoder=Vae_and_Text_Encoders().vae_encoder, 111 | vae_decoder=Vae_and_Text_Encoders().vae_decoder, 112 | text_encoder=Vae_and_Text_Encoders().text_encoder, 113 | tokenizer= self.tokenizer, 114 | unet= self.unet_model, 115 | controlnet=self.controlnet_Model_ort, 116 | scheduler=SchedulersConfig().scheduler(sched_name,model_path), 117 | safety_checker=None, 118 | feature_extractor=None, 119 | requires_safety_checker= False 120 | ) 121 | else: 122 | self.ControlNet_pipe.scheduler= SchedulersConfig().scheduler(sched_name,model_path) 123 | if self.ControlNET_Name!=ControlNET_drop: 124 | self.ControlNet_pipe.controlnet= None 125 | gc.collect() 126 | self.__load_ControlNet_model(ControlNET_drop) 127 | self.ControlNet_pipe.controlnet= self.controlnet_Model_ort 128 | 129 | return self.ControlNet_pipe 130 | 131 | def run_inference(self,prompt,neg_prompt,input_image,width,height,eta,steps,guid,seed,pose_image=None,controlnet_conditioning_scale=1.0): 132 | import numpy as np 133 | rng = np.random.RandomState(int(seed)) 134 | image = self.ControlNet_pipe( 135 | prompt, 136 | input_image, 137 | negative_prompt=neg_prompt, 138 | width = width, 139 | height = height, 140 | num_inference_steps = steps, 141 | guidance_scale=guid, 142 | eta=eta, 143 | num_images_per_prompt=1, 144 | generator=rng, 145 | controlnet_conditioning_scale=controlnet_conditioning_scale 146 | ).images[0] 147 | #Añadir el diccionario 148 | dictio={'prompt':prompt,'neg_prompt':neg_prompt,'steps':steps,'guid':guid,'eta':eta,'strength':controlnet_conditioning_scale,'seed':seed} 149 | return image,dictio 150 | 151 | def create_seeds(self,seed=None,iter=1,same_seeds=False): 152 | self.seeds=seed_generator(seed,iter) 153 | if same_seeds: 154 | for seed in self.seeds: 155 | seed = self.seeds[0] 156 | 157 | def unload_from_memory(self): 158 | self.ControlNet_pipe= None 159 | self.controlnet_Model_ort= None 160 | self.unet_model=None 161 | gc.collect() 162 | 163 | -------------------------------------------------------------------------------- /Engine/engine_common_funcs.py: -------------------------------------------------------------------------------- 1 | 2 | ##Here only common functions who does not call UI variables 3 | import os,re,PIL,numpy 4 | from PIL import Image, PngImagePlugin 5 | 6 | class generator(numpy.random.RandomState): 7 | def __init__(self,seed,generator_type='numpy'): #True = numpy 8 | numpy_or_torch=True if generator_type=='numpy'else False 9 | self.seed=seed 10 | if numpy_or_torch: 11 | super().__init__(self.seed) 12 | else: 13 | import torch 14 | self.torch_gen = torch.Generator() 15 | self.torch_gen.manual_seed(self.seed) 16 | self.randn=self._random 17 | self.randint 18 | 19 | def randint(self,low,max): 20 | import torch 21 | #dato=self.torch_gen.seed() 22 | dato=torch.random(self.torch_gen,low,max) 23 | return dato 24 | 25 | def _random(self,*args): 26 | #def random(self,shape): 27 | shape=tuple(args) 28 | import torch 29 | #torch.manual_seed(self.seed) 30 | #retorno=torch.randn(shape) 31 | retorno=torch.randn(shape, generator=self.torch_gen) 32 | try: 33 | return retorno.numpy() 34 | except: 35 | pass 36 | #finally: 37 | # self=retorno 38 | 39 | 40 | def create_generator(seed,generator_type='numpy'): 41 | import torch 42 | import numpy as np 43 | 44 | if generator_type=='numpy': 45 | generator1=generator(seed,generator_type='numpy') 46 | elif generator_type=='torch': 47 | generator1=generator(seed,generator_type='torch') 48 | else: 49 | print("Generator type does not configured or recognized") 50 | #print(type(generator1)) 51 | return generator1 52 | 53 | def numpy_to_pil(images): 54 | """ 55 | Convert a numpy image or a batch of images to a PIL image. 56 | """ 57 | from PIL import Image 58 | if images.ndim == 3: 59 | images = images[None, ...] 60 | images = (images * 255).round().astype("uint8") 61 | if images.shape[-1] == 1: 62 | # special case for grayscale (single channel) images 63 | pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images] 64 | else: 65 | pil_images = [Image.fromarray(image) for image in images] 66 | return pil_images 67 | 68 | 69 | 70 | def get_next_save_index(output_path): 71 | #output_path=UI_Configuration().output_path 72 | dir_list = os.listdir(output_path) 73 | if len(dir_list): 74 | pattern = re.compile(r"([0-9][0-9][0-9][0-9][0-9][0-9])-([0-9][0-9])\..*") 75 | match_list = [pattern.match(f) for f in dir_list] 76 | next_index = max([int(m[1]) if m else -1 for m in match_list]) + 1 77 | else: 78 | next_index = 0 79 | return next_index 80 | 81 | 82 | def save_image(batch_images,info,next_index,output_path,style=None,save_textfile=False,low_res=False): 83 | #output_path=UI_Configuration().output_path 84 | info_png = f"{info}" 85 | metadata = PngImagePlugin.PngInfo() 86 | metadata.add_text("parameters",info_png) 87 | prompt=info["prompt"] 88 | 89 | style_pre ="" 90 | style_post="" 91 | style_pre_len=0 92 | style_post_len=0 93 | 94 | if style: 95 | styles=style.split("|") 96 | style_pre =styles[0] 97 | style_pre_len=len(style_pre) 98 | style_post=styles[1] 99 | style_post_len=len(style_post) 100 | prompt=prompt[style_pre_len-1:-style_post_len] 101 | """print(f"prompt:{prompt}") 102 | else: 103 | print("No tocado el prompt")""" 104 | 105 | short_prompt1 = prompt.strip("(){}<>:\"/\\|?*\n\t") 106 | short_prompt1 = re.sub(r'[\\/*?:"<>()|\n\t]', "", short_prompt1) 107 | short_prompt = short_prompt1[:49] if len(short_prompt1) > 50 else short_prompt1 108 | os.makedirs(output_path, exist_ok=True) 109 | 110 | i=0 111 | if low_res: 112 | low_res_text="-LowRes" 113 | else: 114 | low_res_text="" 115 | 116 | process_tags=True 117 | if not process_tags: 118 | for image in batch_images: 119 | if image!=None: 120 | image.save(os.path.join(output_path,f"{next_index:06}{low_res_text}-0{i}_{short_prompt}.png",),optimize=True,pnginfo=metadata,) 121 | i+=1 122 | else: 123 | for image in batch_images: 124 | if image!=None: 125 | image.save(os.path.join(output_path,f"{next_index:06}{low_res_text}-0{i}.png",),optimize=True,pnginfo=metadata,) 126 | if save_textfile and (i==0): 127 | #print(f"low_res_text:{low_res_text}-{low_res_text}") 128 | #print(f"low_res_text:{next_index:06}{low_res_text}") 129 | with open(os.path.join(output_path,f"{next_index:06}-{low_res_text}-0{i}.txt"), 'w',encoding='utf8') as txtfile: 130 | txtfile.write(f"{short_prompt1} \nstyle_pre:{style_pre}\nstyle_post:{style_post}") 131 | i+=1 132 | 133 | def PIL_resize_and_crop(input_image: PIL.Image.Image, height: int, width: int): 134 | input_width, input_height = input_image.size 135 | 136 | # nearest neighbor for upscaling 137 | if (input_width * input_height) < (width * height): 138 | resample_type = Image.NEAREST 139 | # lanczos for downscaling 140 | else: 141 | resample_type = Image.LANCZOS 142 | 143 | if height / width > input_height / input_width: 144 | adjust_width = int(input_width * height / input_height) 145 | input_image = input_image.resize((adjust_width, height), 146 | resample=resample_type) 147 | left = (adjust_width - width) // 2 148 | right = left + width 149 | input_image = input_image.crop((left, 0, right, height)) 150 | else: 151 | adjust_height = int(input_height * width / input_width) 152 | input_image = input_image.resize((width, adjust_height), 153 | resample=resample_type) 154 | top = (adjust_height - height) // 2 155 | bottom = top + height 156 | input_image = input_image.crop((0, top, width, bottom)) 157 | return input_image 158 | 159 | 160 | def load_tokenizer_and_config(model_path): 161 | import json 162 | CONFIG_NAME="model_index.json" 163 | config_file=model_path+"/"+CONFIG_NAME 164 | config=None 165 | try: 166 | import json 167 | with open(config_file, "r", encoding="utf-8") as reader: 168 | text = reader.read() 169 | config=json.loads(text) 170 | except: 171 | raise OSError(f"model_index.json not found in {model_path} local folder") 172 | 173 | from transformers import CLIPTokenizer 174 | return CLIPTokenizer.from_pretrained(model_path+"/"+'tokenizer'),config 175 | 176 | def seed_generator(seed,iteration_count): 177 | import numpy as np 178 | # generate seeds for iterations 179 | if seed == "" or seed == None: 180 | rng = np.random.default_rng() 181 | seed = rng.integers(np.iinfo(np.uint32).max) 182 | else: 183 | try: 184 | seed = int(seed) & np.iinfo(np.uint32).max 185 | except ValueError: 186 | seed = hash(seed) & np.iinfo(np.uint32).max 187 | 188 | # use given seed for the first iteration 189 | seeds = np.array([seed], dtype=np.uint32) 190 | 191 | if iteration_count > 1: 192 | seed_seq = np.random.SeedSequence(seed) 193 | seeds = np.concatenate((seeds, seed_seq.generate_state(iteration_count - 1))) 194 | 195 | return seeds 196 | -------------------------------------------------------------------------------- /UI/config_ui_Vae.py: -------------------------------------------------------------------------------- 1 | import gradio as gr 2 | from Engine.General_parameters import running_config 3 | from Engine.General_parameters import VAE_config 4 | 5 | def show_vae_models_configuration(): 6 | VaeConfig=load_vae_preferences__ui() 7 | apply_vae_config_ui(list(VaeConfig.values())) 8 | 9 | 10 | if True: 11 | with gr.Accordion(label="Vae Models Order & Directories",open=False): 12 | gr.Markdown("""Saving disk space\n 13 | Instead of saving duplicates of each VAE (decoder and encoder) files for every model as they are the same, save one instance of generics VAE model\n 14 | into a directory and write down their path.\n 15 | The system will try to apply your 1st option, if not found it go for 2nd and then for 3rd.""") 16 | with gr.Row(): 17 | gr.Markdown("Options for VAE Decoder.",elem_id="title1") 18 | with gr.Row(): 19 | with gr.Column(scale=8): 20 | gr.Markdown("Selected model own VAE decoder.") 21 | with gr.Column(scale=2): 22 | model1_vaedec_order=gr.Slider(1, 3, value=VaeConfig["model1_vaedec_order"], step=1, label="Own model VAE decoder search order", interactive=True) 23 | with gr.Row(): 24 | with gr.Column(scale=8): 25 | model2_vaedec_path=gr.Textbox(label="VAE Decoder model full path",lines=1, value=VaeConfig["model2_vaedec_path"], visible=True, interactive=True) 26 | with gr.Column(scale=2): 27 | model2_vaedec_order=gr.Slider(1, 3, value=VaeConfig["model2_vaedec_order"], step=1, label="This VAE Decoder search order", interactive=True) 28 | with gr.Row(): 29 | with gr.Column(scale=8): 30 | model3_vaedec_path=gr.Textbox(label="VAE Decoder model full path",lines=1, value=VaeConfig["model3_vaedec_path"], visible=True, interactive=True) 31 | with gr.Column(scale=2): 32 | model3_vaedec_order=gr.Slider(1, 3,value=VaeConfig["model3_vaedec_order"], step=1, label="This VAE decoder search order", interactive=True) 33 | with gr.Row(): 34 | gr.Markdown("Options for VAE encoder.",elem_id="title1") 35 | with gr.Row(): 36 | with gr.Column(scale=8): 37 | gr.Markdown("Selected model own VAE encoder.") 38 | with gr.Column(scale=2): 39 | model1_vaeenc_order=gr.Slider(1, 3, value=VaeConfig["model1_vaeenc_order"], step=1, label="Own model VAE Encoder search order", interactive=True) 40 | with gr.Row(): 41 | with gr.Column(scale=8): 42 | model2_vaeenc_path=gr.Textbox(label="VAE Encoder model full path",lines=1, value=VaeConfig["model2_vaeenc_path"], visible=True, interactive=True) 43 | with gr.Column(scale=2): 44 | model2_vaeenc_order=gr.Slider(1, 3, value=VaeConfig["model2_vaeenc_order"], step=1, label="VAE Encoder search order", interactive=True) 45 | with gr.Row(): 46 | with gr.Column(scale=8): 47 | model3_vaeenc_path=gr.Textbox(label="VAE Encoder model full path",lines=1, value=VaeConfig["model3_vaeenc_path"], visible=True, interactive=True) 48 | with gr.Column(scale=2): 49 | model3_vaeenc_order=gr.Slider(1, 3, value=VaeConfig["model3_vaeenc_order"], step=1, label="VAE Encoder search order", interactive=True) 50 | 51 | 52 | save_btn = gr.Button("Apply & Save VAE models config") 53 | load_btn = gr.Button("Load VAE models config") 54 | 55 | all_inputs=[ 56 | model1_vaedec_order,model2_vaedec_order,model2_vaedec_path,model3_vaedec_order,model3_vaedec_path, 57 | model1_vaeenc_order,model2_vaeenc_order,model2_vaeenc_path,model3_vaeenc_order,model3_vaeenc_path] 58 | save_btn.click(fn=save_vae_config_ui, inputs=all_inputs, outputs=None) 59 | load_btn.click(fn=load_vae_preferences__ui2, inputs=None , outputs=all_inputs) 60 | 61 | def load_vae_preferences__ui2(): 62 | return list(load_vae_preferences__ui().values()) 63 | 64 | def load_vae_preferences__ui(): 65 | config=VAE_config() 66 | vaeconfig=config.load_config_from_disk() 67 | #Fast parse,hard-coded, as they are only 3 elements each, for more, do another approach. -recursive funtion 68 | 69 | if vaeconfig[0]=="model": 70 | model1_vaedec_order=1 71 | model2_vaedec_order=2 72 | model2_vaedec_path=vaeconfig[1] 73 | model3_vaedec_order=3 74 | model3_vaedec_path=vaeconfig[2] 75 | elif vaeconfig[1]=="model": 76 | model2_vaedec_order=1 77 | model2_vaedec_path=vaeconfig[0] 78 | model1_vaedec_order=2 79 | model3_vaedec_order=3 80 | model3_vaedec_path=vaeconfig[2] 81 | elif vaeconfig[2]=="model": 82 | model2_vaedec_order=1 83 | model2_vaedec_path=vaeconfig[0] 84 | model3_vaedec_order=2 85 | model3_vaedec_path=vaeconfig[1] 86 | model1_vaedec_order=3 87 | 88 | if vaeconfig[3]=="model": 89 | model1_vaeenc_order=1 90 | model2_vaeenc_order=2 91 | model2_vaeenc_path=vaeconfig[4] 92 | model3_vaeenc_order=3 93 | model3_vaeenc_path=vaeconfig[5] 94 | elif vaeconfig[4]=="model": 95 | model2_vaeenc_order=1 96 | model2_vaeenc_path=vaeconfig[3] 97 | model1_vaeenc_order=2 98 | model3_vaeenc_order=3 99 | model3_vaeenc_path=vaeconfig[5] 100 | elif vaeconfig[5]=="model": 101 | model2_vaeenc_order=1 102 | model2_vaeenc_path=vaeconfig[3] 103 | model3_vaeenc_order=2 104 | model3_vaeenc_path=vaeconfig[4] 105 | model1_vaeenc_order=3 106 | 107 | all_inputs={ 108 | "model1_vaedec_order":model1_vaedec_order, 109 | "model2_vaedec_order":model2_vaedec_order,"model2_vaedec_path":model2_vaedec_path, 110 | "model3_vaedec_order":model3_vaedec_order,"model3_vaedec_path":model3_vaedec_path, 111 | "model1_vaeenc_order":model1_vaeenc_order, 112 | "model2_vaeenc_order":model2_vaeenc_order,"model2_vaeenc_path":model2_vaeenc_path, 113 | "model3_vaeenc_order":model3_vaeenc_order,"model3_vaeenc_path":model3_vaeenc_path} 114 | 115 | return dict(all_inputs) 116 | 117 | 118 | def apply_vae_config_ui(*args): 119 | _save_vae_config_ui(False, *args) 120 | 121 | def save_vae_config_ui(*args): 122 | _save_vae_config_ui(True, *args) 123 | 124 | def _save_vae_config_ui(save=True,*args): 125 | if not save: 126 | args=args[0] #is tupla, select the list of args. 127 | model1_vaedec_order=int(args[0]) 128 | model2_vaedec_order=int(args[1]) 129 | model2_vaedec_path=args[2] 130 | model3_vaedec_order=int(args[3]) 131 | model3_vaedec_path=args[4] 132 | 133 | model1_vaeenc_order=int(args[5]) 134 | model2_vaeenc_order=int(args[6]) 135 | model2_vaeenc_path=args[7] 136 | model3_vaeenc_order=int(args[8]) 137 | model3_vaeenc_path=args[9] 138 | 139 | 140 | 141 | vae_config =[None] * 6 142 | vae_config[model1_vaedec_order-1]="model" 143 | vae_config[model2_vaedec_order-1]=model2_vaedec_path 144 | vae_config[model3_vaedec_order-1]=model3_vaedec_path 145 | vae_config[3+model1_vaeenc_order-1]="model" 146 | vae_config[3+model2_vaeenc_order-1]=model2_vaeenc_path 147 | vae_config[3+model3_vaeenc_order-1]=model3_vaeenc_path 148 | 149 | Running_information= running_config().Running_information 150 | Running_information.update({"Vae_Config":vae_config}) 151 | if save: 152 | VAE_config().save_VAE_config(vae_config) 153 | 154 | 155 | 156 | 157 | -------------------------------------------------------------------------------- /modules/latent_formula_processor.py: -------------------------------------------------------------------------------- 1 | """ 2 | Initialize: 3 | 4 | Return: dict of 3 elements 5 | 1st: a list of tab names where you want it to be included, 6 | 2nd: the area where to be shown in the UI 7 | 3rd: the point where to execute the call 8 | 9 | Additionally you need two funcions: one named "__call__" and another named "__show__" 10 | if no need for an specific function, create one with the return being the same as the imput or a pass function 11 | function __show__: the gradio elements to be shown into the UI component in the specified TAB & AREA, 12 | def show(*args): 13 | pass 14 | function __call__: this will be called to process a dictionay of parameters or a specific parameter, depending on where is to be included 15 | def __call__(*args): 16 | return args 17 | 18 | tabs names:["txt2img","hires"] -- 19 | 20 | 21 | """ 22 | 23 | 24 | 25 | def __init__(*args): 26 | __name__='LatentFormulaProcessor' 27 | print(args[0]) 28 | #here a check of access of initializacion if needed 29 | return { 30 | 'tabs':["hires"], #Añadir tambien a txt2img 31 | 'ui_position':None, #por definir 32 | 'func_processing':None} 33 | 34 | def __call__(*args): 35 | #what to do when the module is called 36 | print("Processing Latent Formula Processor Module") 37 | if args: 38 | datos=args[0] 39 | if type(datos)==dict: 40 | print("LLega dict") 41 | elif type(datos)==str: 42 | print("LLega string") 43 | else: 44 | print("Not recognized input for module Latent Formula Processor %s" % datos) 45 | 46 | return None 47 | 48 | def is_global_ui(): 49 | return False 50 | 51 | def is_global_function(): 52 | return False 53 | 54 | def show(): 55 | show_latent_processor_ui() 56 | 57 | def show_latent_processor_ui(): 58 | import gradio as gr 59 | with gr.Accordion("Latent formula processor module",open=False): 60 | pass 61 | 62 | 63 | def get_ordered_latents(self): 64 | from Engine.General_parameters import running_config 65 | import numpy as np 66 | name=running_config().Running_information["Latent_Name"] 67 | name1= name.split(',') 68 | lista=[0]*len(name1) 69 | for pair in name1: 70 | tupla= pair.split(':') 71 | lista[int(tupla[0])-1]=tupla[1] 72 | #print("Ordered numpys"+str(lista)) 73 | return lista 74 | 75 | def sum_latents(self,latent_list,formula,generator,resultant_latents,iter=0): 76 | #print("Processing formula:"+str(formula)) 77 | subformula_latents= None 78 | while ("(" in formula) or (")" in formula): 79 | #print("Subformula exists") 80 | subformula_startmarks=list([pos for pos, char in enumerate(formula) if char == '(']) 81 | subformula_endmarks=list([pos for pos, char in enumerate(formula) if char == ')']) 82 | 83 | if (len(subformula_endmarks) != len(subformula_startmarks)): 84 | raise Exception("Sorry, Error in formula, check it") 85 | 86 | contador=0 87 | while (len(subformula_startmarks)>contador) and (subformula_startmarks[contador] < subformula_endmarks[0]): 88 | contador+=1 89 | if contador==0: raise Exception("Sorry, Error in formula, check it") 90 | 91 | subformula= formula[(subformula_startmarks[contador-1]+1):subformula_endmarks[0]] 92 | #print(f"subformula:{iter},{subformula}") 93 | previous= formula[0:subformula_startmarks[contador-1]] 94 | posterior=formula[subformula_endmarks[0]+1:] 95 | formula= f"{previous}|{iter}|{posterior}" 96 | iter+=1 97 | subformula_latents = self.sum_latents(latent_list,subformula,generator,resultant_latents,iter) 98 | resultant_latents.append(subformula_latents) 99 | 100 | 101 | # Here we got a plain formula 102 | #print("No subformulas") 103 | result = self.process_simple_formula(latent_list,formula,generator,resultant_latents) 104 | return result 105 | 106 | def process_simple_formula(self,latent_list,formula,generator,resultant_latents): 107 | position=-1 108 | #print("Simple_formula process") 109 | for pos, char in enumerate(formula): 110 | if char in "WwHh": 111 | position=pos 112 | break 113 | if position ==-1 and len(formula)>0: #No operators, single item 114 | result=self.load_latent_file(latent_list,formula,generator,resultant_latents) 115 | else: 116 | previous=formula[0:position] 117 | operator=formula[position] 118 | rest=formula[position+1:] 119 | #print("previous:"+previous) 120 | #print("operator:"+operator) 121 | #print("rest:"+rest) 122 | 123 | result=self.load_latent_file(latent_list,previous,generator,resultant_latents) 124 | result2 = self.process_simple_formula(latent_list,rest,generator,resultant_latents) 125 | 126 | if (operator=='w'): 127 | result = self._sum_latents(result,result2,True) #left & right 128 | elif (operator=='h'): 129 | result = self._sum_latents(result,result2,False) #Up & Down 130 | 131 | return result 132 | 133 | def load_latent_file(self,latent_list,data,generator,resultant_latents): 134 | result = "" 135 | if "|" in data: 136 | lista=data.split("|") 137 | index=int(lista[1]) 138 | result = resultant_latents[index] 139 | #result = "SP:"+resultant_latents[index] 140 | else: 141 | index=int(data) 142 | name=latent_list[int(index)-1] 143 | if "noise" not in name: 144 | print(f"Loading latent(idx:name):{index}:{name}") 145 | result=np.load(f"./latents/{name}") 146 | 147 | """import torch 148 | latents_dtype = result.dtype 149 | noise = generator.randn(*result.shape).astype(latents_dtype) 150 | result = self.hires_pipe.scheduler.add_noise( 151 | torch.from_numpy(result), torch.from_numpy(noise), torch.from_numpy(np.array([1])) 152 | ) 153 | result =result.numpy()""" 154 | else: 155 | noise_size=name.split("noise-")[1].split("x") 156 | print(f"Creating noise block of W/H:{noise_size}") 157 | noise = (0.1)*(generator.random((1,4,int(int(noise_size[1])/8),int(int(noise_size[0])/8))).astype(np.float32)) 158 | result = noise 159 | 160 | return result 161 | 162 | def _sum_latents(self,latent1,latent2,direction): #direction True=horizontal sum(width), False=vertical sum(height) 163 | latent_sum= None 164 | side="" 165 | try: 166 | if direction: 167 | side="Height" 168 | latent_sum = np.concatenate((latent1,latent2),axis=3) #left & right 169 | else: 170 | side="Width" 171 | latent_sum = np.concatenate((latent1,latent2),axis=2) #Up & Down 172 | except: 173 | size1=f"Latent1={(latent1.shape[3]*8)}x{(latent1.shape[2]*8)}" 174 | size2=f"Latent2={(latent2.shape[3]*8)}x{(latent2.shape[2]*8)}" 175 | raise Exception(f"Cannot sum the latents(Width x Height):{size1} and {size2} its {side} must be equal") 176 | return latent_sum 177 | 178 | def get_initial_latent(self, steps,multiplier,generator,strengh): 179 | debug = False 180 | from Engine.General_parameters import running_config 181 | latent_list=self.get_ordered_latents() 182 | formula=running_config().Running_information["Latent_Formula"] 183 | formula=formula.replace(' ', '') 184 | formula=formula.lower() 185 | 186 | loaded_latent=self.sum_latents(latent_list,formula,generator,[]) 187 | 188 | print("Resultant Latent Shape "+"H:"+str(loaded_latent.shape[2]*8)+"x W:"+str(loaded_latent.shape[3]*8)) 189 | 190 | return loaded_latent 191 | -------------------------------------------------------------------------------- /UI/instructp2p_ui.py: -------------------------------------------------------------------------------- 1 | import gradio as gr 2 | import os,re,PIL 3 | 4 | from Engine.General_parameters import Engine_Configuration 5 | from Engine.General_parameters import running_config 6 | from Engine.General_parameters import UI_Configuration as UI_Configuration 7 | from UI import UI_common_funcs as UI_common 8 | from Engine import SchedulersConfig 9 | from Engine import pipelines_engines 10 | 11 | from PIL import Image, PngImagePlugin 12 | 13 | def show_instructp2p_ui(): 14 | ui_config=UI_Configuration() 15 | model_list = UI_common.get_model_list("ip2p") 16 | sched_list = get_schedulers_list() 17 | #sched_list = get_schedulers_list() Modificar a los adecuados 18 | gr.Markdown("Start typing below and then click **Process** to produce the output.Attention: Image is not automatically saved.") 19 | with gr.Row(): 20 | with gr.Column(scale=13, min_width=650): 21 | model_drop = gr.Dropdown(model_list, value=(model_list[0] if len(model_list) > 0 else None), label="model folder", interactive=True) 22 | prompt_t0 = gr.Textbox(value="", lines=2, label="prompt") 23 | sch_t0 = gr.Radio(sched_list, value=sched_list[0], label="scheduler") 24 | image_t0 = gr.Image(source="upload",label="input image", type="pil", elem_id="image_p2p") 25 | iter_t0 = gr.Slider(1, 100, value=1, step=1, label="iteration count", visible=False) 26 | steps_t0 = gr.Slider(1, 300, value=16, step=1, label="steps") 27 | guid_t0 = gr.Slider(0, 50, value=7.5, step=0.1, label="guidance") 28 | height_t0 = gr.Slider(256, 2048, value=512, step=64, label="height") 29 | width_t0 = gr.Slider(256, 2048, value=512, step=64, label="width") 30 | eta_t0 = gr.Slider(0, 1, value=0.0, step=0.01, label="DDIM eta", interactive=True) 31 | seed_t0 = gr.Textbox(value="", max_lines=1, label="seed") 32 | fmt_t0 = gr.Radio(["png", "jpg"], value="png", label="image format", visible=False) 33 | with gr.Column(scale=11, min_width=550): 34 | with gr.Row(): 35 | gen_btn = gr.Button("Process", variant="primary", elem_id="gen_button") 36 | #cancel_btn = gr.Button("Cancel",info="Cancel at end of current iteration",variant="stop", elem_id="gen_button") 37 | memory_btn = gr.Button("Release memory", elem_id="mem_button") 38 | with gr.Row(): 39 | image_out = gr.Gallery(value=None, label="output images") 40 | 41 | with gr.Row(): 42 | status_out = gr.Textbox(value="", label="status") 43 | 44 | 45 | #cancel_btn.click(fn=UI_common.cancel_iteration,inputs=None,outputs=None) 46 | 47 | list_of_All_Parameters=[model_drop,prompt_t0,sch_t0,image_t0,iter_t0,steps_t0,guid_t0,height_t0,width_t0,eta_t0,seed_t0,fmt_t0] 48 | gen_btn.click(fn=generate_click, inputs=list_of_All_Parameters, outputs=[image_out,status_out]) 49 | #sch_t0.change(fn=select_scheduler, inputs=sch_t0, outputs= None) #Atencion cambiar el DDIM ETA si este se activa 50 | memory_btn.click(fn=UI_common.clean_memory_click, inputs=None, outputs=None) 51 | 52 | 53 | def get_model_list(): 54 | model_list = [] 55 | try: 56 | with os.scandir(UI_Configuration().models_dir) as scan_it: 57 | for entry in scan_it: 58 | if entry.is_dir(): 59 | model_list.append(entry.name) 60 | except: 61 | model_list.append("Models directory does not exist, configure it") 62 | return model_list 63 | 64 | def get_schedulers_list(): 65 | sched_config = SchedulersConfig() 66 | sched_list =sched_config.available_schedulers 67 | return sched_list 68 | 69 | def select_scheduler(sched_name,model_path): 70 | return SchedulersConfig().scheduler(sched_name,model_path) 71 | 72 | def generate_click(model_drop,prompt_t0,sch_t0,image_t0,iter_t0,steps_t0,guid_t0,height_t0,width_t0,eta_t0,seed_t0,fmt_t0): 73 | from Engine.pipelines_engines import instruct_p2p_pipe 74 | 75 | Running_information= running_config().Running_information 76 | Running_information.update({"Running":True}) 77 | model_path=ui_config=UI_Configuration().models_dir+"\\"+model_drop 78 | 79 | if (Running_information["model"] != model_drop or Running_information["tab"] != "instruct_p2p"): 80 | UI_common.clean_memory_click() 81 | Running_information.update({"model":model_drop}) 82 | Running_information.update({"tab":"instruct_p2p"}) 83 | 84 | instruct_p2p_pipe().initialize(model_path,sch_t0) 85 | input_image = resize_and_crop(image_t0, height_t0, width_t0) 86 | 87 | # generate seed for iteration 88 | instruct_p2p_pipe().create_seed(seed_t0) 89 | #images= [] 90 | #information=[] 91 | #counter=1 92 | #img_index=get_next_save_index() 93 | images,info = instruct_p2p_pipe().run_inference( 94 | prompt=prompt_t0, 95 | input_image=input_image, 96 | steps=steps_t0, 97 | guid=guid_t0, 98 | eta=eta_t0) 99 | Running_information.update({"Running":False}) 100 | info=dict(info) 101 | info['Sched:']=sch_t0 102 | #information.append(info) 103 | return images,info 104 | 105 | 106 | def get_next_save_index(): 107 | output_path=UI_Configuration().output_path 108 | dir_list = os.listdir(output_path) 109 | if len(dir_list): 110 | pattern = re.compile(r"([0-9][0-9][0-9][0-9][0-9][0-9])-([0-9][0-9])\..*") 111 | match_list = [pattern.match(f) for f in dir_list] 112 | next_index = max([int(m[1]) if m else -1 for m in match_list]) + 1 113 | else: 114 | next_index = 0 115 | return next_index 116 | 117 | 118 | def save_image(batch_images,info,next_index): 119 | output_path=UI_Configuration().output_path 120 | 121 | info_png = f"{info}" 122 | metadata = PngImagePlugin.PngInfo() 123 | metadata.add_text("parameters",info_png) 124 | prompt=info["prompt"] 125 | short_prompt = prompt.strip("<>:\"/\\|?*\n\t") 126 | short_prompt = re.sub(r'[\\/*?:"<>|\n\t]', "", short_prompt) 127 | short_prompt = short_prompt[:49] if len(short_prompt) > 50 else short_prompt 128 | 129 | os.makedirs(output_path, exist_ok=True) 130 | """dir_list = os.listdir(output_path) 131 | if len(dir_list): 132 | pattern = re.compile(r"([0-9][0-9][0-9][0-9][0-9][0-9])-([0-9][0-9])\..*") 133 | match_list = [pattern.match(f) for f in dir_list] 134 | next_index = max([int(m[1]) if m else -1 for m in match_list]) + 1 135 | else: 136 | next_index = 0""" 137 | for image in batch_images: 138 | image.save(os.path.join(output_path,f"{next_index:06}-00.{short_prompt}.png",),optimize=True,pnginfo=metadata,) 139 | 140 | 141 | 142 | 143 | 144 | def resize_and_crop(input_image: PIL.Image.Image, height: int, width: int): 145 | input_width, input_height = input_image.size 146 | 147 | # nearest neighbor for upscaling 148 | if (input_width * input_height) < (width * height): 149 | resample_type = Image.NEAREST 150 | # lanczos for downscaling 151 | else: 152 | resample_type = Image.LANCZOS 153 | 154 | if height / width > input_height / input_width: 155 | adjust_width = int(input_width * height / input_height) 156 | input_image = input_image.resize((adjust_width, height), 157 | resample=resample_type) 158 | left = (adjust_width - width) // 2 159 | right = left + width 160 | input_image = input_image.crop((left, 0, right, height)) 161 | else: 162 | adjust_height = int(input_height * width / input_width) 163 | input_image = input_image.resize((width, adjust_height), 164 | resample=resample_type) 165 | top = (adjust_height - height) // 2 166 | bottom = top + height 167 | input_image = input_image.crop((0, top, width, bottom)) 168 | return input_image -------------------------------------------------------------------------------- /Engine/Pipelines/Vae_and_Text_Encoders.py: -------------------------------------------------------------------------------- 1 | from optimum.onnxruntime.modeling_ort import ORTModel 2 | from optimum.onnxruntime.modeling_diffusion import ORTModelVaeDecoder,ORTModelVaeEncoder,ORTModelTextEncoder 3 | from Engine.General_parameters import Engine_Configuration 4 | 5 | import json 6 | 7 | 8 | class Borg1: 9 | _shared_state = {} 10 | def __init__(self): 11 | self.__dict__ = self._shared_state 12 | 13 | class Vae_and_Text_Encoders(Borg1): 14 | vae_decoder = None 15 | vae_encoder = None 16 | text_encoder = None 17 | 18 | def __init__(self): 19 | Borg1.__init__(self) 20 | 21 | def __str__(self): return json.dumps(self.__dict__) 22 | 23 | def load_vaedecoder(self,model_path,old_version=False): 24 | from Engine.General_parameters import running_config 25 | """if " " in Engine_Configuration().VAEDec_provider: 26 | provider =eval(Engine_Configuration().VAEDec_provider) 27 | else: 28 | provider =Engine_Configuration().VAEDec_provider""" 29 | 30 | provider=Engine_Configuration().VAEDec_provider['provider'] 31 | provider_options=Engine_Configuration().VAEDec_provider['provider_options'] 32 | 33 | running_config=running_config() 34 | import os 35 | if running_config.Running_information["Vae_Config"]: 36 | vae_config=running_config.Running_information["Vae_Config"] 37 | vae_path1= (model_path + "/vae_decoder") if vae_config[0]=="model" else vae_config[0] 38 | vae_path2= (model_path + "/vae_decoder") if vae_config[1]=="model" else vae_config[1] 39 | vae_path3= (model_path + "/vae_decoder") if vae_config[2]=="model" else vae_config[2] 40 | vae_path="" 41 | 42 | if os.path.exists(vae_path1): vae_path= vae_path1 43 | elif os.path.exists(vae_path2): vae_path= vae_path2 44 | elif os.path.exists(vae_path3): vae_path= vae_path3 45 | else: raise Exception("No valid vae decoder path"+vae_path) 46 | 47 | import onnxruntime as ort 48 | sess_options = ort.SessionOptions() 49 | sess_options.log_severity_level=3 50 | 51 | self.vae_decoder = None 52 | print(f"Loading VAE decoder in:{provider}, from {vae_path} with options:{provider_options}" ) 53 | 54 | 55 | 56 | 57 | if old_version: 58 | from diffusers import OnnxRuntimeModel 59 | provider=(str(provider),dict(provider_options)) 60 | self.vae_decoder = OnnxRuntimeModel.from_pretrained(vae_path, provider=provider,sess_options=sess_options) 61 | else: 62 | self.vae_decoder = ORTModel.load_model(vae_path+"/model.onnx", provider,sess_options, provider_options=provider_options) 63 | 64 | return self.vae_decoder #may give problems in the future ( the init of the sd model change it to model, but this is a session) 65 | 66 | """def transform_vadedecodersession_to_model(self,vaedecoder_session): 67 | from optimum.onnxruntime.modeling_diffusion import ORTModelVaeDecoder 68 | return ORTModelVaeDecoder(vaedecoder_session,parent_model=None)""" 69 | 70 | def convert_session_to_model(self,session,model): 71 | if model=="vaedec": 72 | model=ORTModelVaeDecoder(session,None) 73 | elif model=="textenc": 74 | model=ORTModelTextEncoder(session,None) 75 | elif model=="vaeenc": 76 | model=ORTModelVaeEncoder(session,None) 77 | return model 78 | 79 | 80 | def load_vaeencoder(self,model_path,old_version=False): 81 | from Engine.General_parameters import running_config 82 | running_config=running_config() 83 | import os 84 | 85 | vae_config=running_config.Running_information["Vae_Config"] 86 | vae_path1= (model_path + "/vae_encoder") if vae_config[3]=="model" else vae_config[3] 87 | vae_path2= (model_path + "/vae_encoder") if vae_config[4]=="model" else vae_config[4] 88 | vae_path3= (model_path + "/vae_encoder") if vae_config[5]=="model" else vae_config[5] 89 | vae_path="" 90 | 91 | if os.path.exists(vae_path1): vae_path= vae_path1 92 | elif os.path.exists(vae_path2): vae_path= vae_path2 93 | elif os.path.exists(vae_path3): vae_path= vae_path3 94 | else: raise Exception("No valid vae encoder path:"+vae_path) 95 | 96 | """ 97 | if " " in Engine_Configuration().VAEDec_provider: 98 | provider =eval(Engine_Configuration().VAEDec_provider) 99 | else: 100 | provider =Engine_Configuration().VAEDec_provider""" 101 | 102 | provider=Engine_Configuration().VAEEnc_provider['provider'] 103 | provider_options=Engine_Configuration().VAEEnc_provider['provider_options'] 104 | 105 | #vae_path=model_path + "/vae_encoder" 106 | self.vae_encoder = None 107 | 108 | import onnxruntime as ort 109 | sess_options = ort.SessionOptions() 110 | sess_options.log_severity_level=3 111 | print(f"Loading VAE encoder in:{provider}, from {vae_path} with options:{provider_options}") 112 | 113 | if old_version: 114 | from diffusers import OnnxRuntimeModel 115 | provider=(str(provider),dict(provider_options)) 116 | self.vae_encoder = OnnxRuntimeModel.from_pretrained(vae_path, provider=provider,sess_options=sess_options) 117 | else: 118 | self.vae_encoder = ORTModel.load_model(vae_path+"/model.onnx", provider,sess_options, provider_options=provider_options) 119 | 120 | 121 | 122 | 123 | 124 | return self.vae_encoder 125 | 126 | 127 | def load_textencoder(self,model_path,old_version=False): 128 | from Engine.General_parameters import running_config 129 | 130 | import onnxruntime as ort 131 | sess_options = ort.SessionOptions() 132 | sess_options.log_severity_level=3 133 | 134 | provider=Engine_Configuration().TEXTEnc_provider['provider'] 135 | provider_options=Engine_Configuration().TEXTEnc_provider['provider_options'] 136 | 137 | running_config=running_config() 138 | import os 139 | if running_config.Running_information["Textenc_Config"]: 140 | Textenc_Config=running_config.Running_information["Textenc_Config"] 141 | Textenc_path1= (model_path + "/text_encoder") if Textenc_Config[0]=="model" else Textenc_Config[0] 142 | Textenc_path2= (model_path + "/text_encoder") if Textenc_Config[1]=="model" else Textenc_Config[1] 143 | Textenc_path="" 144 | if os.path.exists(Textenc_path1): Textenc_path= Textenc_path1 145 | elif os.path.exists(Textenc_path2): Textenc_path= Textenc_path2 146 | else: raise Exception("No valid Text Encoder path:"+Textenc_path) 147 | 148 | 149 | print(f"Loading TEXT encoder in:{provider} from:{Textenc_path} with options:{provider_options}" ) 150 | self.text_encoder = None 151 | #self.text_encoder = ORTModel.load_model(Textenc_path+"/model.onnx", provider,None,None) 152 | 153 | if old_version: 154 | from diffusers import OnnxRuntimeModel 155 | provider=(str(provider),dict(provider_options)) 156 | #self.text_encoder = ORTModel.load_model(Textenc_path+"/model.onnx" , provider,sess_options,provider_options) 157 | self.text_encoder = OnnxRuntimeModel.from_pretrained(Textenc_path, provider=provider,sess_options=sess_options) 158 | else: 159 | self.text_encoder = ORTModel.load_model(Textenc_path+"/model.onnx" , provider,sess_options,provider_options) 160 | 161 | return self.text_encoder 162 | 163 | def unload_from_memory(self): 164 | import gc 165 | self.vae_decoder = None 166 | self.vae_encoder = None 167 | self.text_encoder = None 168 | gc.collect() 169 | 170 | 171 | -------------------------------------------------------------------------------- /UI/config_ui_TextEncoder.py: -------------------------------------------------------------------------------- 1 | import gradio as gr 2 | from Engine.General_parameters import running_config 3 | from Engine.General_parameters import TextEnc_config 4 | 5 | def show_textenc_models_configuration(): 6 | TextencConfig=load_textenc_preferences__ui() 7 | apply_textenc_config_ui(list(TextencConfig.values())) 8 | 9 | 10 | if True: 11 | with gr.Accordion(label="Text Encoder Models Order & Directories",open=False): 12 | gr.Markdown("""Saving disk space\n 13 | Instead of saving one duplicate for every model as they are usually the same one(two options available only), save one instance of a Text Encoder model\n 14 | into a directory and write down their path.\n 15 | The system will try to apply your 1st option, if not found it will go for 2nd.""") 16 | with gr.Row(): 17 | gr.Markdown("Options for Text Encoder.",elem_id="title1") 18 | with gr.Row(): 19 | with gr.Column(scale=8): 20 | gr.Markdown("Selected model own Text encoder.") 21 | with gr.Column(scale=2): 22 | model1_textenc_order=gr.Slider(1, 5, value=TextencConfig["model1_textenc_order"], step=1, label="Own model Text Encoder search order", interactive=True) 23 | with gr.Row(): 24 | with gr.Column(scale=8): 25 | model2_textenc_path=gr.Textbox(label="Text Encoder clip-slip1",lines=1, value=TextencConfig["model2_textenc_path"], visible=True, interactive=True) 26 | with gr.Column(scale=2): 27 | model2_textenc_order=gr.Slider(1, 5, value=TextencConfig["model2_textenc_order"], step=1, label="This Text Encoder search order", interactive=True) 28 | with gr.Row(): 29 | with gr.Column(scale=8): 30 | model3_textenc_path=gr.Textbox(label="Text Encoder clip-slip2",lines=1, value=TextencConfig["model3_textenc_path"], visible=True, interactive=True) 31 | with gr.Column(scale=2): 32 | model3_textenc_order=gr.Slider(1, 5, value=TextencConfig["model3_textenc_order"], step=1, label="This Text Encoder search order", interactive=True) 33 | with gr.Row(): 34 | with gr.Column(scale=8): 35 | model4_textenc_path=gr.Textbox(label="Text Encoder clip-slip3",lines=1, value=TextencConfig["model4_textenc_path"], visible=True, interactive=True) 36 | with gr.Column(scale=2): 37 | model4_textenc_order=gr.Slider(1, 5, value=TextencConfig["model4_textenc_order"], step=1, label="This Text Encoder search order", interactive=True) 38 | with gr.Row(): 39 | with gr.Column(scale=8): 40 | model5_textenc_path=gr.Textbox(label="Text Encoder clip-slip4",lines=1, value=TextencConfig["model5_textenc_path"], visible=True, interactive=True) 41 | with gr.Column(scale=2): 42 | model5_textenc_order=gr.Slider(1, 5, value=TextencConfig["model5_textenc_order"], step=1, label="This Text Encoder search order", interactive=True) 43 | 44 | save_btn = gr.Button("Apply & Save Text Encoder models config") 45 | load_btn = gr.Button("Load Text Encoder models config") 46 | 47 | all_inputs=[model1_textenc_order, 48 | model2_textenc_order,model2_textenc_path, 49 | model3_textenc_order,model3_textenc_path, 50 | model4_textenc_order,model4_textenc_path, 51 | model5_textenc_order,model5_textenc_path] 52 | save_btn.click(fn=save_textenc_config_ui, inputs=all_inputs, outputs=None) 53 | load_btn.click(fn=load_textenc_preferences__ui2, inputs=None , outputs=all_inputs) 54 | 55 | def load_textenc_preferences__ui2(): 56 | return list(load_textenc_preferences__ui().values()) 57 | 58 | def load_textenc_preferences__ui(): 59 | config=TextEnc_config() 60 | textenc_config=config.load_config_from_disk() 61 | #Fast parse,hard-coded, as they are -*were- only 2 elements each, for more, do another approach. -recursive funtionv(to do) 62 | 63 | 64 | if textenc_config[0]=="model": 65 | model1_textenc_order=1 66 | model2_textenc_order=2 67 | model2_textenc_path=textenc_config[1] 68 | model3_textenc_order=3 69 | model3_textenc_path=textenc_config[2] 70 | model4_textenc_order=4 71 | model4_textenc_path=textenc_config[3] 72 | model5_textenc_order=5 73 | model5_textenc_path=textenc_config[4] 74 | elif textenc_config[1]=="model": 75 | model2_textenc_order=1 76 | model2_textenc_path=textenc_config[0] 77 | model1_textenc_order=2 78 | model3_textenc_order=3 79 | model3_textenc_path=textenc_config[2] 80 | model4_textenc_order=4 81 | model4_textenc_path=textenc_config[3] 82 | model5_textenc_order=5 83 | model5_textenc_path=textenc_config[4] 84 | elif textenc_config[2]=="model": 85 | model2_textenc_order=1 86 | model2_textenc_path=textenc_config[0] 87 | model3_textenc_order=2 88 | model3_textenc_path=textenc_config[1] 89 | model1_textenc_order=3 90 | model4_textenc_order=4 91 | model4_textenc_path=textenc_config[3] 92 | model5_textenc_order=5 93 | model5_textenc_path=textenc_config[4] 94 | elif textenc_config[3]=="model": 95 | model2_textenc_order=1 96 | model2_textenc_path=textenc_config[0] 97 | model3_textenc_order=2 98 | model3_textenc_path=textenc_config[1] 99 | model4_textenc_order=3 100 | model4_textenc_path=textenc_config[2] 101 | model1_textenc_order=4 102 | model5_textenc_order=5 103 | model5_textenc_path=textenc_config[4] 104 | elif textenc_config[4]=="model": 105 | model2_textenc_order=1 106 | model2_textenc_path=textenc_config[0] 107 | model3_textenc_order=2 108 | model3_textenc_path=textenc_config[1] 109 | model4_textenc_order=3 110 | model4_textenc_path=textenc_config[2] 111 | model5_textenc_order=4 112 | model5_textenc_path=textenc_config[3] 113 | model1_textenc_order=5 114 | 115 | 116 | 117 | all_inputs={ 118 | "model1_textenc_order":model1_textenc_order, 119 | "model2_textenc_order":model2_textenc_order, 120 | "model2_textenc_path":model2_textenc_path, 121 | "model3_textenc_order":model3_textenc_order, 122 | "model3_textenc_path":model3_textenc_path, 123 | "model4_textenc_order":model4_textenc_order, 124 | "model4_textenc_path":model4_textenc_path, 125 | "model5_textenc_order":model5_textenc_order, 126 | "model5_textenc_path":model5_textenc_path} 127 | 128 | return dict(all_inputs) 129 | 130 | 131 | def apply_textenc_config_ui(*args): 132 | _save_textenc_config_ui(False, *args) 133 | 134 | def save_textenc_config_ui(*args): 135 | _save_textenc_config_ui(True, *args) 136 | 137 | def _save_textenc_config_ui(save=True,*args): 138 | if not save: 139 | args=args[0] #is tupla, select the list of args. 140 | model1_textenc_order=int(args[0]) 141 | model2_textenc_order=int(args[1]) 142 | model2_textenc_path=args[2] 143 | model3_textenc_order=int(args[3]) 144 | model3_textenc_path=args[4] 145 | model4_textenc_order=int(args[5]) 146 | model4_textenc_path=args[6] 147 | model5_textenc_order=int(args[7]) 148 | model5_textenc_path=args[8] 149 | 150 | 151 | textenc_config =[None] * 5 #number of paths antes eran 2 152 | textenc_config[model1_textenc_order-1]="model" 153 | textenc_config[model2_textenc_order-1]=model2_textenc_path 154 | textenc_config[model3_textenc_order-1]=model3_textenc_path 155 | textenc_config[model4_textenc_order-1]=model4_textenc_path 156 | textenc_config[model5_textenc_order-1]=model5_textenc_path 157 | 158 | 159 | Running_information= running_config().Running_information 160 | Running_information.update({"Textenc_Config":textenc_config}) 161 | if save: 162 | TextEnc_config().save_TextEnc_config(textenc_config) 163 | 164 | 165 | 166 | 167 | -------------------------------------------------------------------------------- /UI/Inpaint_ui.py: -------------------------------------------------------------------------------- 1 | import gradio as gr 2 | import os,re,PIL 3 | 4 | from Engine.General_parameters import running_config 5 | from Engine.General_parameters import UI_Configuration 6 | from UI import UI_common_funcs as UI_common 7 | from Engine import inpaint_pipe 8 | from Engine import SchedulersConfig 9 | #from Engine import pipelines_engines #por mover schedulers_config 10 | #from Engine.General_parameters import Engine_Configuration 11 | 12 | 13 | from PIL import Image, PngImagePlugin 14 | 15 | def show_Inpaint_ui(): 16 | ui_config=UI_Configuration() 17 | model_list = UI_common.get_model_list("inpaint") 18 | sched_list = get_schedulers_list() 19 | gr.Markdown("Start typing below and then click **Process** to produce the output.") 20 | with gr.Row(): 21 | with gr.Column(scale=13, min_width=650): 22 | model_drop = gr.Dropdown(model_list, value=(model_list[0] if len(model_list) > 0 else None), label="model folder", interactive=True) 23 | prompt_t0 = gr.Textbox(value="", lines=2, label="prompt") 24 | neg_prompt_t0 = gr.Textbox(value="", lines=2, label="negative prompt") 25 | sch_t0 = gr.Radio(sched_list, value=sched_list[0], label="scheduler") 26 | #legacy_t0 = gr.Checkbox(value=False, label="legacy inpaint") 27 | 28 | image_t0 = gr.Image( 29 | source="upload", tool="sketch", label="input image", type="pil", elem_id="image_inpaint") 30 | mask_t0 = gr.Image( 31 | source="upload", 32 | label="input mask", 33 | type="pil", 34 | invert_colors=True, 35 | elem_id="mask_inpaint", 36 | ) 37 | with gr.Row(): 38 | iter_t0 = gr.Slider(1, 100, value=1, step=1, label="iteration count") 39 | batch_t0 = gr.Slider(1, 4, value=1, step=1, label="batch size",visible=False) 40 | steps_t0 = gr.Slider(1, 300, value=16, step=1, label="steps") 41 | guid_t0 = gr.Slider(0, 50, value=7.5, step=0.1, label="guidance") 42 | height_t0 = gr.Slider(256, 2048, value=512, step=64, label="height") 43 | width_t0 = gr.Slider(256, 2048, value=512, step=64, label="width") 44 | eta_t0 = gr.Slider(0, 1, value=0.0, step=0.01, label="DDIM eta", interactive=True) 45 | seed_t0 = gr.Textbox(value="", max_lines=1, label="seed") 46 | fmt_t0 = gr.Radio(["png", "jpg"], value="png", label="image format",visible=False) 47 | with gr.Column(scale=11, min_width=550): 48 | with gr.Row(): 49 | gen_btn = gr.Button("Process", variant="primary", elem_id="gen_button") 50 | clear_btn = gr.Button("Cancel",info="Cancel at end of current iteration",variant="stop", elem_id="gen_button") 51 | memory_btn = gr.Button("Release memory", elem_id="mem_button") 52 | 53 | with gr.Row(): 54 | image_out = gr.Gallery(value=None, label="output images") 55 | 56 | with gr.Row(): 57 | status_out = gr.Textbox(value="", label="status") 58 | 59 | 60 | list_of_All_Parameters=[model_drop,prompt_t0,neg_prompt_t0,sch_t0,image_t0,mask_t0,iter_t0,batch_t0,steps_t0,guid_t0,height_t0,width_t0,eta_t0,seed_t0,fmt_t0] 61 | gen_btn.click(fn=generate_click, inputs=list_of_All_Parameters, outputs=[image_out,status_out]) 62 | #sch_t0.change(fn=select_scheduler, inputs=sch_t0, outputs= None) #Atencion cambiar el DDIM ETA si este se activa 63 | memory_btn.click(fn=UI_common.clean_memory_click, inputs=None, outputs=None) 64 | clear_btn.click(fn=UI_common.cancel_iteration,inputs=None,outputs=None) 65 | 66 | 67 | 68 | def gallery_view(images,dict_statuses): 69 | return images[0] 70 | 71 | def get_schedulers_list(): 72 | sched_config = SchedulersConfig() 73 | sched_list =sched_config.available_schedulers 74 | return sched_list 75 | 76 | def select_scheduler(sched_name,model_path): 77 | return SchedulersConfig().scheduler(sched_name,model_path) 78 | 79 | def generate_click(model_drop,prompt_t0,neg_prompt_t0,sch_t0,image_t0,mask_t0,iter_t0,batch_t0,steps_t0,guid_t0,height_t0,width_t0,eta_t0,seed_t0,fmt_t0): 80 | 81 | 82 | Running_information= running_config().Running_information 83 | Running_information.update({"Running":True}) 84 | 85 | #input_image = resize_and_crop(input_image, height_t0, width_t0) 86 | input_image = resize_and_crop(image_t0["image"], height_t0, width_t0) 87 | 88 | if mask_t0 is not None: 89 | print("using uploaded mask") 90 | input_mask = mask_t0.convert("RGB") 91 | input_mask = resize_and_crop(input_mask, height_t0, width_t0) 92 | else: 93 | print("using painted mask") 94 | input_mask = image_t0["mask"].convert("RGB") 95 | input_mask = resize_and_crop(input_mask, height_t0, width_t0) 96 | 97 | if (Running_information["model"] != model_drop or Running_information["tab"] != "inpaint"): 98 | UI_common.clean_memory_click() 99 | Running_information.update({"model":model_drop}) 100 | Running_information.update({"tab":"inpaint"}) 101 | 102 | model_path=ui_config=UI_Configuration().models_dir+"\\"+model_drop 103 | pipe=inpaint_pipe().initialize(model_path,sch_t0) 104 | 105 | inpaint_pipe().create_seeds(seed_t0,iter_t0,False) 106 | images= [] 107 | information=[] 108 | counter=1 109 | img_index=get_next_save_index() 110 | 111 | for seed in inpaint_pipe().seeds: 112 | if running_config().Running_information["cancelled"]: 113 | running_config().Running_information.update({"cancelled":False}) 114 | break 115 | 116 | print(f"Iteration:{counter}/{iter_t0}") 117 | counter+=1 118 | batch_images,info = inpaint_pipe().run_inference( 119 | prompt_t0, 120 | neg_prompt_t0, 121 | input_image, 122 | input_mask, 123 | height_t0, 124 | width_t0, 125 | steps_t0, 126 | guid_t0, 127 | eta_t0, 128 | batch_t0, 129 | seed, 130 | ) 131 | images.extend(batch_images) 132 | info=dict(info) 133 | info['Sched:']=sch_t0 134 | information.append(info) 135 | save_image(batch_images,info,img_index) 136 | img_index+=1 137 | Running_information.update({"Running":False}) 138 | return images,information 139 | 140 | 141 | def get_next_save_index(): 142 | output_path=UI_Configuration().output_path 143 | dir_list = os.listdir(output_path) 144 | if len(dir_list): 145 | pattern = re.compile(r"([0-9][0-9][0-9][0-9][0-9][0-9])-([0-9][0-9])\..*") 146 | match_list = [pattern.match(f) for f in dir_list] 147 | next_index = max([int(m[1]) if m else -1 for m in match_list]) + 1 148 | else: 149 | next_index = 0 150 | return next_index 151 | 152 | 153 | def save_image(batch_images,info,next_index): 154 | output_path=UI_Configuration().output_path 155 | 156 | info_png = f"{info}" 157 | metadata = PngImagePlugin.PngInfo() 158 | metadata.add_text("parameters",info_png) 159 | prompt=info["prompt"] 160 | short_prompt = prompt.strip("<>:\"/\\|?*\n\t") 161 | short_prompt = re.sub(r'[\\/*?:"<>|\n\t]', "", short_prompt) 162 | short_prompt = short_prompt[:49] if len(short_prompt) > 50 else short_prompt 163 | 164 | os.makedirs(output_path, exist_ok=True) 165 | """dir_list = os.listdir(output_path) 166 | if len(dir_list): 167 | pattern = re.compile(r"([0-9][0-9][0-9][0-9][0-9][0-9])-([0-9][0-9])\..*") 168 | match_list = [pattern.match(f) for f in dir_list] 169 | next_index = max([int(m[1]) if m else -1 for m in match_list]) + 1 170 | else: 171 | next_index = 0""" 172 | for image in batch_images: 173 | image.save(os.path.join(output_path,f"{next_index:06}-00.{short_prompt}.png",),optimize=True,pnginfo=metadata,) 174 | 175 | 176 | 177 | def resize_and_crop(input_image: PIL.Image.Image, height: int, width: int): 178 | input_width, input_height = input_image.size 179 | 180 | # nearest neighbor for upscaling 181 | if (input_width * input_height) < (width * height): 182 | resample_type = Image.NEAREST 183 | # lanczos for downscaling 184 | else: 185 | resample_type = Image.LANCZOS 186 | 187 | if height / width > input_height / input_width: 188 | adjust_width = int(input_width * height / input_height) 189 | input_image = input_image.resize((adjust_width, height), 190 | resample=resample_type) 191 | left = (adjust_width - width) // 2 192 | right = left + width 193 | input_image = input_image.crop((left, 0, right, height)) 194 | else: 195 | adjust_height = int(input_height * width / input_width) 196 | input_image = input_image.resize((width, adjust_height), 197 | resample=resample_type) 198 | top = (adjust_height - height) // 2 199 | bottom = top + height 200 | input_image = input_image.crop((0, top, width, bottom)) 201 | return input_image -------------------------------------------------------------------------------- /UI/ControlNet_ui.py: -------------------------------------------------------------------------------- 1 | import gradio as gr 2 | import os,gc,re,PIL 3 | 4 | from Engine.General_parameters import Engine_Configuration 5 | from Engine.General_parameters import running_config 6 | from Engine.General_parameters import UI_Configuration 7 | from Engine.General_parameters import ControlNet_config 8 | from UI import UI_common_funcs as UI_common 9 | from Engine import SchedulersConfig 10 | #from Engine import pipelines_engines 11 | from Engine import ControlNet_pipe 12 | 13 | from PIL import Image, PngImagePlugin 14 | 15 | 16 | def show_ControlNet_ui(): 17 | ui_config=UI_Configuration() 18 | model_list = UI_common.get_model_list("controlnet") 19 | sched_list = get_schedulers_list() 20 | controlnet_models= dict(ControlNet_config().available_controlnet_models()) 21 | gr.Markdown("Start typing below and then click **Process** to produce the output.") 22 | with gr.Row(): 23 | with gr.Column(scale=13, min_width=650): 24 | model_drop = gr.Dropdown(model_list, value=(model_list[0] if len(model_list) > 0 else None), label="model folder", interactive=True) 25 | with gr.Row(): 26 | ControlNET_drop = gr.Dropdown(controlnet_models.keys(),value=next(iter(controlnet_models)), label="ControNET Type", interactive=True) 27 | reload_controlnet_btn = gr.Button("Reload ControlNet model list") 28 | 29 | prompt_t0 = gr.Textbox(value="", lines=2, label="prompt") 30 | neg_prompt_t0 = gr.Textbox(value="", lines=2, label="negative prompt") 31 | sch_t0 = gr.Radio(sched_list, value=sched_list[0], label="scheduler", interactive=True) 32 | 33 | image_t0 = gr.Image(source="upload",label="input image", type="pil", elem_id="image_inpaint", visible=True) 34 | #pose_image_t0 = gr.Image(source="upload",label="ControlNET Image",type="pil") 35 | 36 | with gr.Row(): 37 | iter_t0 = gr.Slider(1, 100, value=1, step=1, label="iteration count") 38 | batch_t0 = gr.Slider(1, 4, value=1, step=1, label="batch size",visible=False) 39 | steps_t0 = gr.Slider(1, 300, value=16, step=1, label="steps") 40 | guid_t0 = gr.Slider(0, 50, value=7.5, step=0.1, label="guidance") 41 | controlnet_conditioning_scale_t0= gr.Slider(0, 1, value=0.5, step=0.1, label="controlnet_conditioning_scale") 42 | with gr.Row(): 43 | height_t0 = gr.Slider(256, 2048, value=512, step=64, label="height") 44 | width_t0 = gr.Slider(256, 2048, value=512, step=64, label="width") 45 | eta_t0 = gr.Slider(0, 1, value=0.0, step=0.01, label="DDIM eta", interactive=False) 46 | seed_t0 = gr.Textbox(value="", max_lines=1, label="seed") 47 | #fmt_t0 = gr.Radio(["png", "jpg"], value="png", label="image format") 48 | with gr.Column(scale=11, min_width=550): 49 | with gr.Row(): 50 | gen_btn = gr.Button("Process Full from Image", variant="primary", elem_id="gen_button") 51 | #gen_btn2 = gr.Button("Process to get pose image", variant="primary", elem_id="gen_button") 52 | #gen_btn3 = gr.Button("Process from pose image", variant="primary", elem_id="gen_button") 53 | with gr.Row(): 54 | clear_btn = gr.Button("Cancel",info="Cancel at end of current iteration",variant="stop", elem_id="gen_button") 55 | memory_btn = gr.Button("Release memory", elem_id="mem_button") 56 | with gr.Row(): 57 | image_out = gr.Gallery(value=None, label="output images") 58 | with gr.Row(): 59 | status_out = gr.Textbox(value="", label="status") 60 | 61 | list_of_All_Parameters=[model_drop,prompt_t0,neg_prompt_t0,sch_t0,image_t0,iter_t0,batch_t0,steps_t0,guid_t0,height_t0,width_t0,eta_t0,seed_t0,ControlNET_drop,controlnet_conditioning_scale_t0] 62 | gen_btn.click(fn=generate_click, inputs=list_of_All_Parameters, outputs=[image_out,status_out]) 63 | #gen_btn1.click(fn=generate_click, inputs=list_of_All_Parameters, outputs=[pose_image_t0,status_out]) 64 | #gen_btn2.click(fn=generate_click, inputs=list_of_All_Parameters, outputs=[image_out,status_out]) 65 | #sch_t0.change(fn=select_scheduler, inputs=sch_t0, outputs= None) #Atencion cambiar el DDIM ETA si este se activa 66 | memory_btn.click(fn=UI_common.clean_memory_click, inputs=None, outputs=None) 67 | clear_btn.click(fn=UI_common.cancel_iteration,inputs=None,outputs=None) 68 | reload_controlnet_btn.click(fn=ReloadControlNet, inputs=None , outputs=ControlNET_drop) 69 | 70 | 71 | def ReloadControlNet(): 72 | ControlNet_config().load_config_from_disk() 73 | controlnet_models= dict(ControlNet_config().available_controlnet_models()) 74 | controlnet_models_keys= list(controlnet_models.keys()) 75 | return gr.Dropdown.update(choices=controlnet_models_keys, value=controlnet_models_keys[0]) 76 | 77 | 78 | def get_schedulers_list(): 79 | sched_config = SchedulersConfig() 80 | sched_list =sched_config.schedulers_controlnet_list() 81 | return sched_list 82 | 83 | def select_scheduler(sched_name,model_path): 84 | return SchedulersConfig().scheduler(sched_name,model_path) 85 | 86 | 87 | def generate_click(model_drop,prompt_t0,neg_prompt_t0,sch_t0,image_t0,iter_t0,batch_t0,steps_t0,guid_t0,height_t0,width_t0,eta_t0,seed_t0,ControlNET_drop,controlnet_conditioning_scale): 88 | 89 | Running_information= running_config().Running_information 90 | Running_information.update({"Running":True}) 91 | 92 | input_image = resize_and_crop(image_t0, height_t0, width_t0) 93 | #pose_image_t0 = resize_and_crop(pose_image_t0, height_t0, width_t0) 94 | 95 | model_path=UI_Configuration().models_dir+"\\"+model_drop 96 | 97 | if (Running_information["model"] != model_drop or Running_information["tab"] != "controlnet"): 98 | UI_common.clean_memory_click() 99 | Running_information.update({"model":model_drop}) 100 | Running_information.update({"tab":"controlnet"}) 101 | ControlNet_pipe().initialize(model_path,sch_t0,ControlNET_drop) 102 | 103 | 104 | ControlNet_pipe().create_seeds(seed_t0,iter_t0,False) 105 | images= [] 106 | information=[] 107 | counter=1 108 | img_index=get_next_save_index() 109 | 110 | for seed in ControlNet_pipe().seeds: 111 | if running_config().Running_information["cancelled"]: 112 | running_config().Running_information.update({"cancelled":False}) 113 | break 114 | 115 | print(f"Iteration:{counter}/{iter_t0}") 116 | counter+=1 117 | batch_images,info = ControlNet_pipe().run_inference( 118 | prompt_t0, 119 | neg_prompt=neg_prompt_t0, 120 | input_image=input_image, 121 | #pose_image=pose_image_t0, 122 | height=height_t0, 123 | width=width_t0, 124 | steps=steps_t0, 125 | guid=guid_t0, 126 | eta=eta_t0, 127 | controlnet_conditioning_scale=controlnet_conditioning_scale, 128 | seed=seed) 129 | 130 | images.extend([batch_images]) 131 | info=dict(info) 132 | info['Sched:']=sch_t0 133 | info['CnetModel:']=ControlNET_drop 134 | information.append(info) 135 | save_image([batch_images],info,img_index) 136 | img_index+=1 137 | Running_information.update({"Running":False}) 138 | return images,information 139 | 140 | 141 | def get_next_save_index(): 142 | output_path=UI_Configuration().output_path 143 | dir_list = os.listdir(output_path) 144 | if len(dir_list): 145 | pattern = re.compile(r"([0-9][0-9][0-9][0-9][0-9][0-9])-([0-9][0-9])\..*") 146 | match_list = [pattern.match(f) for f in dir_list] 147 | next_index = max([int(m[1]) if m else -1 for m in match_list]) + 1 148 | else: 149 | next_index = 0 150 | return next_index 151 | 152 | 153 | def save_image(batch_images,info,next_index): 154 | output_path=UI_Configuration().output_path 155 | 156 | info_png = f"{info}" 157 | metadata = PngImagePlugin.PngInfo() 158 | metadata.add_text("parameters",info_png) 159 | prompt=info["prompt"] 160 | short_prompt = prompt.strip("<>:\"/\\|?*\n\t") 161 | short_prompt = re.sub(r'[\\/*?:"<>|\n\t]', "", short_prompt) 162 | short_prompt = short_prompt[:49] if len(short_prompt) > 50 else short_prompt 163 | 164 | os.makedirs(output_path, exist_ok=True) 165 | 166 | for image in batch_images: 167 | image.save(os.path.join(output_path,f"{next_index:06}-00.{short_prompt}.png",),optimize=True,pnginfo=metadata,) 168 | 169 | 170 | def resize_and_crop(input_image: PIL.Image.Image, height: int, width: int): 171 | input_width, input_height = input_image.size 172 | 173 | # nearest neighbor for upscaling 174 | if (input_width * input_height) < (width * height): 175 | resample_type = Image.NEAREST 176 | # lanczos for downscaling 177 | else: 178 | resample_type = Image.LANCZOS 179 | 180 | if height / width > input_height / input_width: 181 | adjust_width = int(input_width * height / input_height) 182 | input_image = input_image.resize((adjust_width, height), 183 | resample=resample_type) 184 | left = (adjust_width - width) // 2 185 | right = left + width 186 | input_image = input_image.crop((left, 0, right, height)) 187 | else: 188 | adjust_height = int(input_height * width / input_width) 189 | input_image = input_image.resize((width, adjust_height), 190 | resample=resample_type) 191 | top = (adjust_height - height) // 2 192 | bottom = top + height 193 | input_image = input_image.crop((0, top, width, bottom)) 194 | return input_image -------------------------------------------------------------------------------- /UI/config_ui_engine.py: -------------------------------------------------------------------------------- 1 | import gradio as gr 2 | import onnxruntime as ort 3 | from Engine.General_parameters import Engine_Configuration 4 | 5 | global debug 6 | debug = False 7 | 8 | def show_providers_configuration(): 9 | with gr.Blocks(title="ONNX Difussers UI-2") as test: 10 | own_code="{'device_id': 1}" 11 | Engine_Config=Engine_Configuration() 12 | with gr.Accordion(label="Select which provider will load&execute each process/pipeline",open=False): 13 | gr.Markdown("Device provider options(dict), i.e. a python dict {'device_id': 1} will select the 2nd device provider, for 2 GPUs available") 14 | with gr.Row(): 15 | MAINPipe_Select=gr.Radio(ort.get_available_providers(), 16 | label="MAINPipes provider", info="Where to load the main pipes",value=Engine_Config.MAINPipe_provider['provider']) 17 | MAINPipe_own_code=gr.Textbox(label="Pipeline provider",info="Device provider options(dict)",lines=1, value=Engine_Config.MAINPipe_provider['provider_options'], visible=True, interactive=True) 18 | #MAINPipe_Select.change(fn=Generic_Select_OwnCode,inputs=MAINPipe_Select,outputs=MAINPipe_own_code) 19 | with gr.Row(): 20 | Sched_Select=gr.Radio(ort.get_available_providers(), label="Scheduler provider", info="Where to load the schedulers",value=Engine_Config.Scheduler_provider['provider']) 21 | Sched_own_code=gr.Textbox(label="Pipeline provider",info="Device provider options(dict)",lines=1, value=Engine_Config.Scheduler_provider['provider_options'], visible=True) 22 | #Sched_Select.change(fn=Generic_Select_OwnCode,inputs=Sched_Select,outputs=Sched_own_code) 23 | with gr.Row(): 24 | ControlNet_Select=gr.Radio(ort.get_available_providers(), label="ControlNet provider", info="Where to load ControlNet",value=Engine_Config.ControlNet_provider['provider']) 25 | ControlNet_own_code=gr.Textbox(label="Pipeline provider",info="Device provider options(dict)",lines=1, value=Engine_Config.ControlNet_provider['provider_options'], visible=True) 26 | #ControlNet_Select.change(fn=Generic_Select_OwnCode,inputs=ControlNet_Select,outputs=ControlNet_own_code) 27 | with gr.Row(): 28 | VAEDec_Select=gr.Radio(ort.get_available_providers(), label="VAEDecoder provider", info="Where to load the VAE Decoder",value=Engine_Config.VAEDec_provider['provider']) 29 | VAEDec_own_code=gr.Textbox(label="Pipeline provider",info="Device provider options(dict)",lines=1, value=Engine_Config.VAEDec_provider['provider_options'], visible=True) 30 | #VAEDec_Select.change(fn=Generic_Select_OwnCode,inputs=VAEDec_Select,outputs=VAEDec_own_code) 31 | with gr.Row(): 32 | VAEEnc_Select=gr.Radio(ort.get_available_providers(), label="VAEEncoder provider", info="Where to load the VAE Encoder",value=Engine_Config.VAEEnc_provider['provider']) 33 | VAEEnc_own_code=gr.Textbox(label="Pipeline provider",info="Device provider options(dict)",lines=1, value=Engine_Config.VAEEnc_provider['provider_options'], visible=True) 34 | #VAEDec_Select.change(fn=Generic_Select_OwnCode,inputs=VAEDec_Select,outputs=VAEDec_own_code) 35 | with gr.Row(): 36 | TEXTEnc_Select=gr.Radio(ort.get_available_providers(), label="TEXTEncoder provider", info="Where to load the Text Encoder",value=Engine_Config.TEXTEnc_provider['provider']) 37 | TEXTEnc_own_code=gr.Textbox(label="Pipeline provider",info="Device provider options(dict)",lines=1, value=Engine_Config.TEXTEnc_provider['provider_options'], visible=True) 38 | #TEXTEnc_Select.change(fn=Generic_Select_OwnCode,inputs=TEXTEnc_Select,outputs=TEXTEnc_own_code) 39 | with gr.Row(): 40 | LowResPipe_Select=gr.Radio(ort.get_available_providers(), label="Low Res Pipeline provider", info="Where to load the LowRes Model in HiRes pipeline",value=Engine_Config.LowResPipe_provider['provider']) 41 | LowResPipe_own_code=gr.Textbox(label="Pipeline provider",info="Device provider options(dict)",lines=1, value=Engine_Config.LowResPipe_provider['provider_options'], visible=True) 42 | #TEXTEnc_Select.change(fn=Generic_Select_OwnCode,inputs=TEXTEnc_Select,outputs=TEXTEnc_own_code) 43 | 44 | with gr.Row(): 45 | DeepDanbooru_Select=gr.Radio(ort.get_available_providers(), label="DeepDanbooru provider", info="Where to load DeepDanbooru queries",value=Engine_Config.DeepDanBooru_provider) 46 | gr.Markdown("DeepDanbooru provider. If already loaded for a query, unload it to apply changes") 47 | with gr.Row(): 48 | apply_btn = gr.Button("Apply providers config", variant="primary") 49 | apply_btn.click(fn=apply_config, inputs=[MAINPipe_Select, MAINPipe_own_code,Sched_Select,Sched_own_code,ControlNet_Select,ControlNet_own_code, VAEDec_Select,VAEDec_own_code,TEXTEnc_Select,TEXTEnc_own_code,DeepDanbooru_Select,VAEEnc_Select,VAEEnc_own_code,LowResPipe_Select,LowResPipe_own_code] , outputs=None) 50 | save_btn = gr.Button("Save to config file", elem_id="test_button") 51 | save_btn.click(fn=save_config_disk, inputs=None, outputs=None) 52 | load_btn = gr.Button("Load from config file", elem_id="test_button") 53 | load_btn.click(fn=load_config_disk, inputs=None, outputs=[MAINPipe_Select, MAINPipe_own_code,Sched_Select,Sched_own_code,ControlNet_Select,ControlNet_own_code, VAEDec_Select,VAEDec_own_code,TEXTEnc_Select,TEXTEnc_own_code,DeepDanbooru_Select, VAEEnc_Select,VAEEnc_own_code,LowResPipe_Select,LowResPipe_own_code]) 54 | 55 | 56 | def apply_config(MAINPipe_Select, MAINPipe_own_code,Sched_Select,Sched_own_code,ControlNet_Select,ControlNet_own_code, 57 | VAEDec_Select,VAEDec_own_code,TEXTEnc_Select,TEXTEnc_own_code,DeepDanbooru_Select,VAEEnc_Select,VAEEnc_own_code,LowResPipe_Select,LowResPipe_own_code): 58 | Engine_Config=Engine_Configuration() 59 | Engine_Config.SetProvider('MAINPipe_provider',get_provider_code('MAINPipe',MAINPipe_Select,MAINPipe_own_code)) 60 | Engine_Config.SetProvider('Scheduler_provider',get_provider_code('Schedulers',Sched_Select,Sched_own_code)) 61 | Engine_Config.SetProvider('ControlNet_provider',get_provider_code('ControlNet',ControlNet_Select,ControlNet_own_code)) 62 | Engine_Config.SetProvider('VAEDec_provider',get_provider_code('VaeDecoder',VAEDec_Select,VAEDec_own_code)) 63 | Engine_Config.SetProvider('VAEEnc_provider',get_provider_code('VaeEncoder',VAEEnc_Select,VAEEnc_own_code)) #Asegurarnos que existe antes de activar la linea 64 | Engine_Config.SetProvider('TEXTEnc_provider',get_provider_code('TextEncoder',TEXTEnc_Select,TEXTEnc_own_code)) 65 | Engine_Config.SetProvider('LowResPipe_provider',get_provider_code('LowResPipe',LowResPipe_Select,LowResPipe_own_code)) 66 | 67 | 68 | Engine_Config.DeepDanBooru_provider=DeepDanbooru_Select 69 | global debug 70 | if debug: 71 | print("Applied a new running config:") 72 | print_current_singleton_data(Engine_Configuration()) 73 | 74 | def print_current_singleton_data(Engine_Config): 75 | print(Engine_Config.MAINPipe_provider) 76 | print(Engine_Config.Scheduler_provider) 77 | print(Engine_Config.ControlNet_provider) 78 | print(Engine_Config.VAEDec_provider) 79 | print(Engine_Config.TEXTEnc_provider) 80 | print(Engine_Config.DeepDanBooru_provider) 81 | 82 | def save_config_disk(): 83 | global debug 84 | if debug: 85 | print("Saving to disk the Running Singleton content(Maybe not the UI content):") 86 | print_current_singleton_data(Engine_Configuration()) 87 | Engine_Configuration().save_config_json() 88 | 89 | def load_config_disk(): 90 | global debug 91 | if debug: 92 | print("Retrieving from disk the Singleton content:") 93 | print_current_singleton_data(Engine_Configuration()) 94 | Engine_Configuration().load_config_json() 95 | #Update gr.radio values for the UI to be able to update it. 96 | loaded=Engine_Configuration() 97 | test = loaded.MAINPipe_provider['provider'] 98 | test2 = loaded.MAINPipe_provider['provider_options'] 99 | test3 = loaded.Scheduler_provider['provider'] 100 | test4 = loaded.Scheduler_provider['provider_options'] 101 | test5 = loaded.ControlNet_provider['provider'] 102 | test6 = loaded.ControlNet_provider['provider_options'] 103 | test7 = loaded.VAEDec_provider['provider'] 104 | test8 = loaded.VAEDec_provider['provider_options'] 105 | test9 = loaded.TEXTEnc_provider['provider'] 106 | test10 = loaded.TEXTEnc_provider['provider_options'] 107 | test11 = loaded.DeepDanBooru_provider 108 | test12 = loaded.VAEEnc_provider['provider'] 109 | test13 = loaded.VAEEnc_provider['provider_options'] 110 | test14 = loaded.LowResPipe_provider['provider'] 111 | test15 = loaded.LowResPipe_provider['provider_options'] 112 | return test,test2,test3,test4,test5,test6,test7,test8,test9,test10,test11,test12,test13,test14,test15 113 | 114 | 115 | def get_provider_code_old(Selection,OwnCode): #Esta funcion es eliminable si no hay que modificar el retorno 116 | if Selection =="Own_code": 117 | return OwnCode 118 | else: 119 | return str(Selection) 120 | 121 | 122 | def get_provider_code(pipename,Selection,OwnCode): 123 | try: 124 | provider_options=eval(OwnCode) 125 | except: 126 | provider_options="" 127 | 128 | if type(provider_options)==dict: 129 | OwnCode=provider_options 130 | else: 131 | OwnCode=None 132 | 133 | #return {pipename+'_provider':Selection, pipename+'_provider_options':provider_options} 134 | return {'provider':Selection, 'provider_options':provider_options} 135 | 136 | def Generic_Select_OwnCode(Btn_Select): #No utilizada, para version antigua que permitia otras opciones 137 | if Btn_Select == "Own_code": 138 | return gr.Textbox.update(visible=True) 139 | else: 140 | return gr.Textbox.update(visible=False) 141 | 142 | -------------------------------------------------------------------------------- /UI/Img2Img_ui.py: -------------------------------------------------------------------------------- 1 | import gradio as gr 2 | import os,gc,re,PIL 3 | from PIL import Image, PngImagePlugin 4 | 5 | from Engine.General_parameters import Engine_Configuration 6 | from Engine.General_parameters import running_config 7 | from Engine.General_parameters import UI_Configuration 8 | from UI import UI_common_funcs as UI_common 9 | from Engine import SchedulersConfig 10 | #from Engine import pipelines_engines 11 | 12 | 13 | 14 | def show_Img2Img_ui(): 15 | ui_config=UI_Configuration() 16 | model_list = UI_common.get_model_list("txt2img") 17 | sched_list = get_schedulers_list() 18 | gr.Markdown("Start typing below and then click **Process** to produce the output.") 19 | with gr.Row(): 20 | with gr.Column(scale=13, min_width=650): 21 | model_drop = gr.Dropdown(model_list, value=(model_list[0] if len(model_list) > 0 else None), label="model folder", interactive=True) 22 | prompt_t0 = gr.Textbox(value="", lines=2, label="prompt") 23 | neg_prompt_t0 = gr.Textbox(value="", lines=2, label="negative prompt") 24 | sch_t0 = gr.Radio(sched_list, value=sched_list[0], label="scheduler") 25 | 26 | image_t0 = gr.Image( 27 | tool="sketch", label="input image", type="pil", elem_id="image2image") 28 | #source="upload", tool="sketch", label="input image", type="pil", elem_id="image2image") 29 | 30 | with gr.Row(): 31 | iter_t0 = gr.Slider(1, 100, value=1, step=1, label="iteration count") 32 | batch_t0 = gr.Slider(1, 4, value=1, step=1, label="batch size",visible=False) 33 | steps_t0 = gr.Slider(1, 300, value=16, step=1, label="steps") 34 | guid_t0 = gr.Slider(0, 50, value=7.5, step=0.1, label="guidance") 35 | strengh_t0 = gr.Slider(0, 1, value=0.75, step=0.05, label="Strengh") 36 | height_t0 = gr.Slider(256, 2048, value=512, step=64, label="Height") 37 | width_t0 = gr.Slider(256, 2048, value=512, step=64, label="Width") 38 | eta_t0 = gr.Slider(0, 1, value=0.0, step=0.01, label="DDIM eta", interactive=True) 39 | seed_t0 = gr.Textbox(value="", max_lines=1, label="seed") 40 | fmt_t0 = gr.Radio(["png", "jpg"], value="png", label="image format",visible=False) 41 | with gr.Column(scale=11, min_width=550): 42 | with gr.Row(): 43 | gen_btn = gr.Button("Process", variant="primary", elem_id="gen_button") 44 | cancel_btn = gr.Button("Cancel",info="Cancel at end of current iteration",variant="stop", elem_id="gen_button") 45 | memory_btn = gr.Button("Release memory", elem_id="mem_button") 46 | with gr.Row(): 47 | image_out = gr.Gallery(value=None, label="output images") 48 | with gr.Row(): 49 | status_out = gr.Textbox(value="", label="status") 50 | 51 | 52 | list_of_All_Parameters=[model_drop,prompt_t0,neg_prompt_t0,sch_t0,image_t0,iter_t0,batch_t0,steps_t0,guid_t0,height_t0,width_t0,eta_t0,seed_t0,fmt_t0,strengh_t0] 53 | gen_btn.click(fn=generate_click, inputs=list_of_All_Parameters, outputs=[image_out,status_out]) 54 | #sch_t0.change(fn=select_scheduler, inputs=sch_t0, outputs= None) #Atencion cambiar el DDIM ETA si este se activa 55 | memory_btn.click(fn=UI_common.clean_memory_click, inputs=None, outputs=None) 56 | cancel_btn.click(fn=UI_common.cancel_iteration,inputs=None,outputs=None) 57 | 58 | 59 | 60 | def gallery_view(images,dict_statuses): 61 | return images[0] 62 | 63 | 64 | def get_schedulers_list(): 65 | sched_config = SchedulersConfig() 66 | sched_list =sched_config.available_schedulers 67 | return sched_list 68 | 69 | def select_scheduler(sched_name,model_path): 70 | return SchedulersConfig().scheduler(sched_name,model_path) 71 | 72 | def generate_click(model_drop,prompt_t0,neg_prompt_t0,sch_t0,image_t0,iter_t0,batch_t0,steps_t0,guid_t0,height_t0,width_t0,eta_t0,seed_t0,fmt_t0,strengh_t0): 73 | #from Engine.pipelines_engines import img2img_pipe 74 | from Engine import img2img_pipe 75 | Running_information= running_config().Running_information 76 | Running_information.update({"Running":True}) 77 | 78 | input_image = resize_and_crop(image_t0["image"], height_t0, width_t0) 79 | 80 | if (Running_information["model"] != model_drop or Running_information["tab"] != "img2img"): 81 | #if (Running_information["model"] != model_drop or Running_information["tab"] != "txt2img"): 82 | UI_common.clean_memory_click() 83 | Running_information.update({"model":model_drop}) 84 | Running_information.update({"tab":"img2img"}) 85 | 86 | model_path=ui_config=UI_Configuration().models_dir+"\\"+model_drop 87 | pipe=img2img_pipe().initialize(model_path,sch_t0) 88 | img2img_pipe().create_seeds(seed_t0,iter_t0,False) 89 | images= [] 90 | information=[] 91 | counter=1 92 | img_index=get_next_save_index() 93 | 94 | for seed in img2img_pipe().seeds: 95 | if running_config().Running_information["cancelled"]: 96 | running_config().Running_information.update({"cancelled":False}) 97 | break 98 | 99 | print(f"Iteration:{counter}/{iter_t0}") 100 | counter+=1 101 | loopback =False 102 | if loopback is True: 103 | try: 104 | loopback_image 105 | except UnboundLocalError: 106 | loopback_image = None 107 | 108 | if loopback_image is not None: 109 | batch_images,info = img2img_pipe().run_inference( 110 | prompt_t0, 111 | negative_prompt=neg_prompt_t0, 112 | image=loopback_image, #Check this data, 113 | num_inference_steps=steps_t0, 114 | guidance_scale=guid_t0, 115 | eta=eta_t0, 116 | strength=strengh_t0, 117 | num_images_per_prompt=batch_t0, 118 | seed=seed).images 119 | elif loopback_image is None: 120 | batch_images,info = img2img_pipe().run_inference( 121 | prompt_t0, 122 | negative_prompt=neg_prompt_t0, 123 | image=input_image, 124 | strength=strengh_t0, 125 | num_inference_steps=steps_t0, 126 | guidance_scale=guid_t0, 127 | eta=eta_t0, 128 | num_images_per_prompt=batch_t0, 129 | seed=seed).images 130 | elif loopback is False: 131 | batch_images,info = img2img_pipe().run_inference( 132 | prompt_t0, 133 | neg_prompt=neg_prompt_t0, 134 | init_image=input_image, 135 | strength=strengh_t0, 136 | steps=steps_t0, 137 | guid=guid_t0, 138 | eta=eta_t0, 139 | batch=batch_t0, 140 | seed=seed) #Cambio gen por su seed.cuidado falta eta pero podria hacer falta 141 | 142 | images.extend(batch_images) 143 | info=dict(info) 144 | info['Sched:']=sch_t0 145 | information.append(info) 146 | save_image(batch_images,info,img_index) 147 | img_index+=1 148 | Running_information.update({"Running":False}) 149 | return images,information 150 | 151 | 152 | def get_next_save_index(): 153 | output_path=UI_Configuration().output_path 154 | dir_list = os.listdir(output_path) 155 | if len(dir_list): 156 | pattern = re.compile(r"([0-9][0-9][0-9][0-9][0-9][0-9])-([0-9][0-9])\..*") 157 | match_list = [pattern.match(f) for f in dir_list] 158 | next_index = max([int(m[1]) if m else -1 for m in match_list]) + 1 159 | else: 160 | next_index = 0 161 | return next_index 162 | 163 | 164 | def save_image(batch_images,info,next_index): 165 | output_path=UI_Configuration().output_path 166 | 167 | info_png = f"{info}" 168 | metadata = PngImagePlugin.PngInfo() 169 | metadata.add_text("parameters",info_png) 170 | prompt=info["Img2ImgPrompt"] 171 | short_prompt = prompt.strip("<>:\"/\\|?*\n\t") 172 | short_prompt = re.sub(r'[\\/*?:"<>|\n\t]', "", short_prompt) 173 | short_prompt = short_prompt[:49] if len(short_prompt) > 50 else short_prompt 174 | 175 | os.makedirs(output_path, exist_ok=True) 176 | """dir_list = os.listdir(output_path) 177 | if len(dir_list): 178 | pattern = re.compile(r"([0-9][0-9][0-9][0-9][0-9][0-9])-([0-9][0-9])\..*") 179 | match_list = [pattern.match(f) for f in dir_list] 180 | next_index = max([int(m[1]) if m else -1 for m in match_list]) + 1 181 | else: 182 | next_index = 0""" 183 | for image in batch_images: 184 | image.save(os.path.join(output_path,f"{next_index:06}-00.{short_prompt}.png",),optimize=True,pnginfo=metadata,) 185 | 186 | 187 | def resize_and_crop(input_image: PIL.Image.Image, height: int, width: int): 188 | input_width, input_height = input_image.size 189 | 190 | # nearest neighbor for upscaling 191 | if (input_width * input_height) < (width * height): 192 | resample_type = Image.NEAREST 193 | # lanczos for downscaling 194 | else: 195 | resample_type = Image.LANCZOS 196 | 197 | if height / width > input_height / input_width: 198 | adjust_width = int(input_width * height / input_height) 199 | input_image = input_image.resize((adjust_width, height), 200 | resample=resample_type) 201 | left = (adjust_width - width) // 2 202 | right = left + width 203 | input_image = input_image.crop((left, 0, right, height)) 204 | else: 205 | adjust_height = int(input_height * width / input_width) 206 | input_image = input_image.resize((width, adjust_height), 207 | resample=resample_type) 208 | top = (adjust_height - height) // 2 209 | bottom = top + height 210 | input_image = input_image.crop((0, top, width, bottom)) 211 | return input_image -------------------------------------------------------------------------------- /IMG-Strip.py: -------------------------------------------------------------------------------- 1 | import gradio as gr 2 | from pathlib import Path 3 | from PIL import Image 4 | #from DIS_ISNET.models import * 5 | 6 | 7 | Metadata_value="Created by @" 8 | 9 | def init_ui(): 10 | downloads_path = str(Path.home() / "Downloads") 11 | with gr.Blocks(title="Additional Tools") as img_strip: 12 | with gr.Tab(label="Remove & Extract Metadata") as Metadata_tab: 13 | image_in = gr.Image(label="input image", type="pil", elem_id="image_init") 14 | save_dir = gr.Textbox(label="Save to Dir",value=downloads_path, lines=1) 15 | old_metadata = gr.Textbox(label="Existing metadata",value="", lines=4) 16 | new_metadata = gr.Textbox(label="Metadata to overwrite",value=Metadata_value, lines=4) 17 | threshold = gr.Slider(0, 255, value=20, step=1, label="Mask threshold", interactive=True) 18 | extract_metadata_btn = gr.Button("Extract Metadata", variant="primary") 19 | process_btn = gr.Button("Process Metadata & Save Image", variant="primary") 20 | convert_btn = gr.Button("Extract mask of main objects") 21 | image_out = gr.Image(label="Object") 22 | image_out2 = gr.Image(label="Detected Mask") 23 | 24 | 25 | process_btn.click(fn=delete_metadata_saveimg, inputs=[image_in,save_dir,new_metadata] , outputs=None) 26 | extract_metadata_btn.click(fn=extract_metadata, inputs=image_in , outputs=old_metadata) 27 | convert_btn.click(fn=test_fn, inputs=[image_in,threshold] , outputs=[image_out2,image_out]) 28 | return img_strip 29 | 30 | 31 | 32 | def test_fn(image,threshold): 33 | #import torch.onnx 34 | #import torch 35 | #import torch.nn.functional as F 36 | #from torchvision.transforms.functional import normalize 37 | 38 | import numpy as np 39 | import cv2 40 | import onnxruntime 41 | 42 | 43 | input_size=[1024,1024] #size of model input width X height 44 | 45 | input_image = np.array(image.convert('RGB')) 46 | #input_image = cv2.cvtColor(input_image, cv2.COLOR_RGB2BGR) 47 | input_image_resized_np=cv2.resize(input_image,input_size) 48 | 49 | input_shape=input_image.shape[0:2] #size of input image height X width 50 | 51 | if len(input_image_resized_np.shape) < 3: 52 | input_image_resized_np = input_image_resized_np[:, :, np.newaxis] 53 | 54 | input_image_resized_np=input_image_resized_np.transpose(2,0,1)#.astype(np.uint8) 55 | input_image_resized_np=np.expand_dims(input_image_resized_np, 0) 56 | input_image_resized_np=input_image_resized_np/255.0 57 | 58 | min_val = np.min(input_image_resized_np) 59 | max_val = np.max(input_image_resized_np) 60 | input_image_resized_np = (input_image_resized_np - min_val) / (max_val - min_val) #Normalization 61 | 62 | 63 | #session = onnxruntime.InferenceSession("model_isnet.onnx_fp16.onnx",providers=['DmlExecutionProvider', 'CPUExecutionProvider']) 64 | #results = session.run(["output"], {"image": input_image_resized_np.astype(np.float16}) 65 | session = onnxruntime.InferenceSession("./modules/dis-isnet/model_isnet.onnx",providers=['CUDAExecutionProvider','DmlExecutionProvider', 'CPUExecutionProvider']) 66 | results = session.run(["output"], {"image": input_image_resized_np.astype(np.float32)}) 67 | 68 | result=results[0] 69 | #result=torch.from_numpy(result) 70 | #result=torch.squeeze(F.upsample(result,input_shape,mode='bilinear'),0) #from pth inference converted to numpy 71 | 72 | result=np.squeeze(result) 73 | max=result.max() 74 | min=result.min() 75 | result =(result-min)/(max-min) #Normalization 76 | result = (result *255).astype(np.uint8) 77 | 78 | 79 | result_3channels=cv2.merge([result,result,result]) #create 3 duplicated channels for each color to be considered as image 80 | result_3channels=cv2.resize(result_3channels,(input_shape[1],input_shape[0])) 81 | result,_,_=cv2.split(result_3channels) 82 | 83 | result=np.expand_dims(result, axis=2) 84 | 85 | 86 | #result =result.transpose(1,2,0).astype(np.uint8) 87 | _,result = cv2.threshold(result, threshold, 255, cv2.THRESH_BINARY) #mirar el threshold, esta en 10, pero inicial en 120 88 | #final=cv2.merge([input_image,result]) 89 | result1=cv2.bitwise_or(input_image,input_image,mask=result) ##pensar si dejar fondo(input_image) o borrarlo (result1) 90 | #result=cv2.merge([input_image,result]) 91 | result=cv2.merge([result1,result]) 92 | #result=cv2.cvtColor(result, cv2.COLOR_BGR2RGB) 93 | 94 | image_applied_mask=Image.fromarray(result) 95 | 96 | return result_3channels,image_applied_mask 97 | 98 | #return result_3channels,final 99 | 100 | 101 | def test_fn2(image): #pruebas de cambio de pth a onnx 102 | 103 | import torch.onnx 104 | import torch 105 | import numpy as np 106 | import cv2 107 | #import torch.nn as nn 108 | #from torchvision import models 109 | import torch.nn.functional as F 110 | from torchvision.transforms.functional import normalize 111 | 112 | input_size=[1024,1024] 113 | 114 | 115 | model_path = "./DIS_ISNET/isnet.pth" 116 | #model=ISNetDIS() 117 | model.load_state_dict(torch.load(model_path,map_location=torch.device('cpu'))) 118 | #print("Dict") 119 | #print(model) 120 | new_image = np.array(image.convert('RGB')) 121 | im=new_image 122 | new_image = cv2.cvtColor(new_image, cv2.COLOR_RGB2BGR) 123 | 124 | #im= np.array(image).copy() #convert to skimage opposite(to PIL) : Image.fromarray(image) 125 | 126 | with torch.no_grad(): 127 | if len(im.shape) < 3: 128 | im = im[:, :, np.newaxis] 129 | im_shp=im.shape[0:2] 130 | im_tensor = torch.tensor(im, dtype=torch.float32).permute(2,0,1) 131 | im_tensor = F.upsample(torch.unsqueeze(im_tensor,0), input_size, mode="bilinear").type(torch.uint8) 132 | image = torch.divide(im_tensor,255.0) 133 | image = normalize(image,[0.5,0.5,0.5],[1.0,1.0,1.0]) 134 | 135 | if torch.cuda.is_available(): 136 | image=image.cuda() 137 | 138 | import onnxruntime 139 | session = onnxruntime.InferenceSession("model_isnet.onnx",providers=['CUDAExecutionProvider','DmlExecutionProvider', 'CPUExecutionProvider']) 140 | results = session.run(["output"], {"image": image.numpy()}) 141 | 142 | result=results[0] 143 | result=torch.from_numpy(result) 144 | #print(type(result)) 145 | #print(result.shape) 146 | 147 | #result=model(image) 148 | #print(type(result)) 149 | #print(type(result[0][0])) 150 | #print(result.shape) 151 | 152 | #result=torch.squeeze(F.upsample(result[0][0],im_shp,mode='bilinear'),0) 153 | result=torch.squeeze(F.upsample(result,im_shp,mode='bilinear'),0) 154 | ma = torch.max(result) 155 | mi = torch.min(result) 156 | result = (result-mi)/(ma-mi) 157 | result=(result*255) 158 | 159 | 160 | result=result.permute(1,2,0).cpu().data.numpy().astype(np.uint8) 161 | result2=cv2.merge([result,result,result]) 162 | result = cv2.cvtColor(result2, cv2.COLOR_RGB2GRAY) 163 | _,result = cv2.threshold(result, 120, 255, cv2.THRESH_BINARY) 164 | #print(result.shape) 165 | 166 | """model.eval() 167 | x = torch.randn(1, 3, 1024, 1024, requires_grad=True) 168 | x.to('cpu') 169 | #input_names = ["batch,channels,height,width"]"input_ids" 170 | #input_names = ['b','c','height','width'] 171 | input_names = ['image'] 172 | output_names = ["output0"] 173 | #torch_out = model(x) 174 | #torch.onnx.export(model, x, "model_isnet.onnx", export_params=True, opset_version=11, do_constant_folding=True, input_names = ['image'], output_names = ['output'], dynamic_axes = {'image' : {2: 'height',3:'width'}, 'output': {2: 'height',3:'width'}}) 175 | #torch.onnx.export(model, x, "model_isnet.onnx", export_params=True, opset_version=11, do_constant_folding=True, input_names = ['input'], output_names = ['output']) 176 | torch.onnx.export(model, x, "model_isnet.onnx", export_params=True, opset_version=11, do_constant_folding=True, input_names = ['image'], output_names = ['output'])""" 177 | 178 | 179 | 180 | result=cv2.bitwise_or(new_image,new_image,mask=result) 181 | result=cv2.cvtColor(result, cv2.COLOR_BGR2RGB) 182 | image2=Image.fromarray(result) 183 | """ 184 | 185 | import onnxruntime 186 | session = onnxruntime.InferenceSession("model_isnet.onnx",providers=['DmlExecutionProvider', 'CPUExecutionProvider']) 187 | session.get_modelmeta() 188 | first_input_name = session.get_inputs()[0].name 189 | first_output_name = session.get_outputs()[0].name 190 | print("Modelo onnx") 191 | print(first_input_name) 192 | print(first_output_name) 193 | 194 | 195 | print("The model expects input shape: ", session.get_inputs()[0].shape) 196 | print("The shape of the Image is: ", mantener.shape) 197 | results = session.run(["output"], {"image": mantener.numpy()}) 198 | 199 | result=results[0] 200 | print(type(result))""" 201 | """import onnx 202 | from onnxconverter_common import float16 203 | model = onnx.load("model_isnet.onnx") 204 | model_fp16 = float16.convert_float_to_float16(model) 205 | onnx.save(model_fp16, "model_isnet.onnx_fp16.onnx")""" 206 | 207 | return result2,image2 208 | 209 | 210 | def extract_metadata(image): 211 | #metadata=list(image.info.values()) 212 | metadata=image.info 213 | return metadata 214 | 215 | 216 | def delete_metadata_saveimg(image,save_dir, new_metadata=""): 217 | from PIL import PngImagePlugin 218 | #exifdata = eval(image.info['parameters']) 219 | #exifdata = (image.info['parameters']) 220 | exifdata =str(image.info) 221 | import hashlib 222 | new_name= hashlib.sha256(exifdata.encode()).hexdigest()[0:10] 223 | 224 | metadata = PngImagePlugin.PngInfo() 225 | metadata.add_text("ImageInfo",f"{new_metadata}") 226 | 227 | image.save(f"{save_dir}/{new_name}.png",optimize=True,pnginfo=metadata) 228 | 229 | 230 | img_strip =init_ui() 231 | img_strip.launch(server_port=8080) 232 | 233 | 234 | -------------------------------------------------------------------------------- /Engine/General_parameters.py: -------------------------------------------------------------------------------- 1 | 2 | # Singleton/BorgSingleton.py 3 | # Alex Martelli's 'Borg' 4 | # https://python-3-patterns-idioms-test.readthedocs.io/en/latest/Singleton.html 5 | import json 6 | 7 | class Borg: 8 | _shared_state = {} 9 | def __init__(self): 10 | self.__dict__ = self._shared_state 11 | 12 | class Borg1: 13 | _shared_state = {} 14 | def __init__(self): 15 | self.__dict__ = self._shared_state 16 | 17 | class Borg2: 18 | _shared_state = {} 19 | def __init__(self): 20 | self.__dict__ = self._shared_state 21 | 22 | class Borg3: 23 | _shared_state = {} 24 | def __init__(self): 25 | self.__dict__ = self._shared_state 26 | class Borg4: 27 | _shared_state = {} 28 | def __init__(self): 29 | self.__dict__ = self._shared_state 30 | 31 | class Borg5: 32 | _shared_state = {} 33 | def __init__(self): 34 | self.__dict__ = self._shared_state 35 | 36 | 37 | class Engine_Configuration(Borg): 38 | MAINPipe_provider="Not Selected" 39 | Scheduler_provider="Not Selected" 40 | ControlNet_provider="Not Selected" 41 | VAEDec_provider="Not Selected" 42 | VAEEnc_provider="Not Selected" 43 | TEXTEnc_provider="Not Selected" 44 | DeepDanBooru_provider="Not Selected" 45 | LowResPipe_provider="Not Selected" 46 | 47 | #MAINPipe_provider,Scheduler_provider,ControlNet_provider,VAEDec_provider,TEXTEnc_provider 48 | def __init__(self): 49 | Borg.__init__(self) 50 | 51 | def __str__(self): return json.dumps(self.__dict__) 52 | 53 | def SetProvider(self,pipeline,diccionary): 54 | setattr(self, pipeline, diccionary) 55 | print(getattr(self, pipeline)) 56 | 57 | def save_config_json(self): 58 | jsonStr = json.dumps(self.__dict__) 59 | #print(type(jsonStr)) 60 | with open("./Engine/config_files/EngineConfig_nueva.json", "w") as outfile: 61 | outfile.write(jsonStr) 62 | print("Saving information for pipelines providers") 63 | print(jsonStr) 64 | return jsonStr 65 | 66 | def load_default_values(self): 67 | print("Loading default provider values:CPU") 68 | self.MAINPipe_provider={'provider':"CPUExecutionProvider",'provider_options':{"device_id": 0} } 69 | self.Scheduler_provider={'provider':"CPUExecutionProvider",'provider_options':{"device_id": 0} } 70 | self.ControlNet_provider={'provider':"CPUExecutionProvider",'provider_options':{"device_id": 0} } 71 | self.VAEDec_provider={'provider':"CPUExecutionProvider",'provider_options':{"device_id": 0} } 72 | self.VAEEnc_provider={'provider':"CPUExecutionProvider",'provider_options':{"device_id": 0} } 73 | self.TEXTEnc_provider={'provider':"CPUExecutionProvider",'provider_options':{"device_id": 0} } 74 | self.LowResPipe_provider={'provider':"CPUExecutionProvider",'provider_options':{"device_id": 0} } 75 | self.DeepDanBooru_provider="CPUExecutionProvider" 76 | 77 | def load_config_json(self): 78 | try: 79 | with open('./Engine/config_files/EngineConfig_nueva.json', 'r') as openfile: 80 | jsonStr = json.load(openfile) 81 | self.MAINPipe_provider = jsonStr["MAINPipe_provider"] 82 | self.Scheduler_provider = jsonStr["Scheduler_provider"] 83 | self.ControlNet_provider = jsonStr["ControlNet_provider"] 84 | self.VAEDec_provider = jsonStr["VAEDec_provider"] 85 | self.VAEEnc_provider = jsonStr["VAEEnc_provider"] 86 | self.TEXTEnc_provider = jsonStr["TEXTEnc_provider"] 87 | self.DeepDanBooru_provider = jsonStr["DeepDanBooru_provider"] 88 | self.LowResPipe_provider = jsonStr["LowResPipe_provider"] 89 | except OSError: 90 | self.load_default_values() 91 | return self 92 | 93 | class UI_Configuration(Borg1): 94 | __loaded= False 95 | models_dir="" 96 | output_path = "" 97 | wildcards_activated=True 98 | Txt2img_Tab = None 99 | InPaint_Tab = None 100 | Img2Img_Tab = None 101 | InstructP2P_Tab = None 102 | ControlNet_Tab = None 103 | Tools_Tab = None 104 | Advanced_Config = None 105 | GradioPort = 7860 106 | 107 | def __init__(self): 108 | Borg1.__init__(self) 109 | if not self.__loaded: 110 | self.load_config() 111 | 112 | def __str__(self): return json.dumps(self.__dict__) 113 | 114 | def __load_default_values(self): 115 | import os 116 | self.models_dir=os.getcwd()+"\\models" 117 | self.output_path=os.getcwd()+"\\output" 118 | self.Txt2img_Tab = True 119 | self.InPaint_Tab = True 120 | self.Img2Img_Tab = True 121 | self.Tools_Tab = True 122 | self.InstructP2P_Tab = True 123 | self.ControlNet_Tab = True 124 | self.Advanced_Config = True 125 | self.GradioPort = 7860 126 | 127 | def save_config_json(self): 128 | jsonStr = json.dumps(self.__dict__) 129 | with open("./Engine/config_files/UIConfig.json", "w") as outfile: 130 | outfile.write(jsonStr) 131 | print("Saving UI Config") 132 | print(jsonStr) 133 | return jsonStr 134 | 135 | def __load_config_json(self): 136 | try: 137 | with open('./Engine/config_files/UIConfig.json', 'r') as openfile: 138 | jsonStr = json.load(openfile) 139 | self.models_dir = jsonStr["models_dir"] 140 | self.output_path= jsonStr["output_path"] 141 | self.Txt2img_Tab = jsonStr["Txt2img_Tab"] 142 | self.InPaint_Tab = jsonStr["InPaint_Tab"] 143 | self.Img2Img_Tab = jsonStr["Img2Img_Tab"] 144 | self.InstructP2P_Tab = jsonStr["InstructP2P_Tab"] 145 | self.ControlNet_Tab = int(jsonStr["ControlNet_Tab"]) 146 | self.Tools_Tab = jsonStr["Tools_Tab"] 147 | self.Advanced_Config = jsonStr["Advanced_Config"] 148 | self.GradioPort = int(jsonStr["GradioPort"]) 149 | 150 | except OSError: 151 | self.__load_default_values() 152 | return self 153 | 154 | def load_config(self): 155 | self.__load_config_json() 156 | self.__loaded=True 157 | 158 | 159 | class running_config(Borg2): 160 | Running_information= dict({"loaded":False}) 161 | 162 | def __init__(self): 163 | Borg2.__init__(self) 164 | if not self.Running_information["loaded"]==True: 165 | self.Running_information.update({"loaded":True}) 166 | 167 | def __str__(self): return json.dumps(self.__dict__) 168 | 169 | class VAE_config(Borg4): 170 | config = None 171 | def __init__(self): 172 | Borg4.__init__(self) 173 | if self.config == None: 174 | self.config = self.__load_VAE_config() 175 | 176 | def __load_VAE_config(self): 177 | import json 178 | standard_config = None 179 | try: 180 | with open('./Engine/config_files/VAE_Preferences.json', 'r') as openfile: 181 | jsonStr = json.load(openfile) 182 | vae_config = list(jsonStr) 183 | except OSError: 184 | standard_config = ["model"]*6 185 | self.config=vae_config 186 | return vae_config 187 | 188 | def save_VAE_config(self,vae_config): 189 | print("Saving VAE Config") 190 | 191 | import json 192 | json_data=jsonstr=json.dumps(vae_config) 193 | with open("./Engine/config_files/VAE_Preferences.json", "w") as outfile: 194 | outfile.write(json_data) 195 | 196 | def load_config_from_disk(self): 197 | self.config = self.__load_VAE_config() 198 | return self.config 199 | 200 | class TextEnc_config(Borg5): 201 | config = None 202 | def __init__(self): 203 | Borg5.__init__(self) 204 | if self.config == None: 205 | self.config = self.__load_TextEnc_config() 206 | 207 | def __load_TextEnc_config(self): 208 | import json 209 | TextEnc_config = None 210 | try: 211 | with open('./Engine/config_files/TextEnc_Preferences.json', 'r') as openfile: 212 | jsonStr = json.load(openfile) 213 | TextEnc_config = list(jsonStr) 214 | except OSError: 215 | TextEnc_config = ["model"]*5 216 | self.config=TextEnc_config 217 | return TextEnc_config 218 | 219 | def save_TextEnc_config(self,TextEnc_config): 220 | import json 221 | json_data=json.dumps(TextEnc_config) 222 | with open("./Engine/config_files/TextEnc_Preferences.json", "w") as outfile: 223 | outfile.write(json_data) 224 | 225 | def load_config_from_disk(self): 226 | self.config = self.__load_TextEnc_config() 227 | return self.config 228 | 229 | class ControlNet_config(Borg3): 230 | config = None 231 | def __init__(self): 232 | Borg3.__init__(self) 233 | if self.config == None: 234 | self.config = self.__load_controlnet_config() 235 | 236 | def __str__(self): return json.dumps(self.__dict__) 237 | 238 | def __load_controlnet_config(self): 239 | import json 240 | standard_config = None 241 | try: 242 | with open('./Engine/config_files/ControlNet.json', 'r') as openfile: 243 | jsonStr = json.load(openfile) 244 | standard_config = dict(jsonStr) 245 | except OSError: 246 | standard_config = { 247 | "canny_active":False, 248 | "canny_path":"", 249 | "depth_active":False, 250 | "depth_path":"", 251 | "hed_active":False, 252 | "hed_path":"", 253 | "mlsd_active":False, 254 | "mlsd_path":"", 255 | "normal_active":False, 256 | "normal_path":"", 257 | "openpose_active":False, 258 | "openpose_path":"", 259 | "seg_active":False, 260 | "seg_path":"", 261 | } 262 | return standard_config 263 | 264 | def load_config_from_disk(self): 265 | self.config = self.__load_controlnet_config() 266 | self.available_controlnet_models() 267 | 268 | def save_controlnet_config(self,controlnet_config): 269 | print("Saving ControlNet models config") 270 | 271 | import json 272 | json_data=jsonstr=json.dumps(controlnet_config) 273 | with open("./Engine/config_files/ControlNet.json", "w") as outfile: 274 | outfile.write(json_data) 275 | 276 | def available_controlnet_models(self): 277 | available=[] 278 | for key, value in self.config.items(): 279 | if "active" in key and value == True: 280 | model=key.split('_')[0] 281 | available.append((model,self.config[model+"_path"])) 282 | return available #a list of tuples (model, path) 283 | 284 | 285 | -------------------------------------------------------------------------------- /Engine/txt2img_pipe_sub.py: -------------------------------------------------------------------------------- 1 | ''' 2 | This code is copied from original pipeline, adding only the break option for callback''' 3 | 4 | import inspect 5 | from typing import Callable, List, Optional, Union 6 | 7 | import numpy as np 8 | import torch 9 | from transformers import CLIPImageProcessor, CLIPTokenizer 10 | 11 | from diffusers.configuration_utils import FrozenDict 12 | from diffusers.schedulers import DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler 13 | from diffusers.utils import deprecate, logging 14 | from diffusers.pipelines.onnx_utils import ORT_TO_NP_TYPE, OnnxRuntimeModel 15 | from diffusers.pipelines.pipeline_utils import DiffusionPipeline 16 | from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput 17 | 18 | def __call__( 19 | self, 20 | prompt: Union[str, List[str]] = None, 21 | height: Optional[int] = 512, 22 | width: Optional[int] = 512, 23 | num_inference_steps: Optional[int] = 50, 24 | guidance_scale: Optional[float] = 7.5, 25 | negative_prompt: Optional[Union[str, List[str]]] = None, 26 | num_images_per_prompt: Optional[int] = 1, 27 | eta: Optional[float] = 0.0, 28 | generator: Optional[np.random.RandomState] = None, 29 | latents: Optional[np.ndarray] = None, 30 | prompt_embeds: Optional[np.ndarray] = None, 31 | negative_prompt_embeds: Optional[np.ndarray] = None, 32 | output_type: Optional[str] = "pil", 33 | return_dict: bool = True, 34 | callback: Optional[Callable[[int, int, np.ndarray], None]] = None, 35 | callback_steps: int = 1, 36 | ): 37 | r""" 38 | Function invoked when calling the pipeline for generation. 39 | 40 | Args: 41 | prompt (`str` or `List[str]`, *optional*): 42 | The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. 43 | instead. 44 | image (`PIL.Image.Image` or List[`PIL.Image.Image`] or `torch.FloatTensor`): 45 | `Image`, or tensor representing an image batch which will be upscaled. * 46 | num_inference_steps (`int`, *optional*, defaults to 50): 47 | The number of denoising steps. More denoising steps usually lead to a higher quality image at the 48 | expense of slower inference. 49 | guidance_scale (`float`, *optional*, defaults to 7.5): 50 | Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). 51 | `guidance_scale` is defined as `w` of equation 2. of [Imagen 52 | Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > 53 | 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, 54 | usually at the expense of lower image quality. 55 | negative_prompt (`str` or `List[str]`, *optional*): 56 | The prompt or prompts not to guide the image generation. If not defined, one has to pass 57 | `negative_prompt_embeds`. instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` 58 | is less than `1`). 59 | num_images_per_prompt (`int`, *optional*, defaults to 1): 60 | The number of images to generate per prompt. 61 | eta (`float`, *optional*, defaults to 0.0): 62 | Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to 63 | [`schedulers.DDIMScheduler`], will be ignored for others. 64 | generator (`np.random.RandomState`, *optional*): 65 | One or a list of [numpy generator(s)](TODO) to make generation deterministic. 66 | latents (`np.ndarray`, *optional*): 67 | Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image 68 | generation. Can be used to tweak the same generation with different prompts. If not provided, a latents 69 | tensor will ge generated by sampling using the supplied random `generator`. 70 | prompt_embeds (`np.ndarray`, *optional*): 71 | Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not 72 | provided, text embeddings will be generated from `prompt` input argument. 73 | negative_prompt_embeds (`np.ndarray`, *optional*): 74 | Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt 75 | weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input 76 | argument. 77 | output_type (`str`, *optional*, defaults to `"pil"`): 78 | The output format of the generate image. Choose between 79 | [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. 80 | return_dict (`bool`, *optional*, defaults to `True`): 81 | Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a 82 | plain tuple. 83 | callback (`Callable`, *optional*): 84 | A function that will be called every `callback_steps` steps during inference. The function will be 85 | called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`. 86 | callback_steps (`int`, *optional*, defaults to 1): 87 | The frequency at which the `callback` function will be called. If not specified, the callback will be 88 | called at every step. 89 | 90 | Returns: 91 | [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`: 92 | [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple. 93 | When returning a tuple, the first element is a list with the generated images, and the second element is a 94 | list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work" 95 | (nsfw) content, according to the `safety_checker`. 96 | """ 97 | #print("Using substituted pipe to allow inferences cancel") 98 | # check inputs. Raise error if not correct 99 | self.check_inputs( 100 | prompt, height, width, callback_steps, negative_prompt, prompt_embeds, negative_prompt_embeds 101 | ) 102 | 103 | # define call parameters 104 | if prompt is not None and isinstance(prompt, str): 105 | batch_size = 1 106 | elif prompt is not None and isinstance(prompt, list): 107 | batch_size = len(prompt) 108 | else: 109 | batch_size = prompt_embeds.shape[0] 110 | 111 | if generator is None: 112 | generator = np.random 113 | 114 | # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) 115 | # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` 116 | # corresponds to doing no classifier free guidance. 117 | do_classifier_free_guidance = guidance_scale > 1.0 118 | 119 | prompt_embeds = self._encode_prompt( 120 | prompt, 121 | num_images_per_prompt, 122 | do_classifier_free_guidance, 123 | negative_prompt, 124 | prompt_embeds=prompt_embeds, 125 | negative_prompt_embeds=negative_prompt_embeds, 126 | ) 127 | 128 | # get the initial random noise unless the user supplied it 129 | latents_dtype = prompt_embeds.dtype 130 | latents_shape = (batch_size * num_images_per_prompt, 4, height // 8, width // 8) 131 | if latents is None: 132 | latents = generator.randn(*latents_shape).astype(latents_dtype) 133 | elif latents.shape != latents_shape: 134 | raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {latents_shape}") 135 | 136 | # set timesteps 137 | self.scheduler.set_timesteps(num_inference_steps) 138 | 139 | latents = latents * np.float64(self.scheduler.init_noise_sigma) 140 | 141 | # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature 142 | # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers. 143 | # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502 144 | # and should be between [0, 1] 145 | accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys()) 146 | extra_step_kwargs = {} 147 | if accepts_eta: 148 | extra_step_kwargs["eta"] = eta 149 | 150 | timestep_dtype = next( 151 | (input.type for input in self.unet.model.get_inputs() if input.name == "timestep"), "tensor(float)" 152 | ) 153 | timestep_dtype = ORT_TO_NP_TYPE[timestep_dtype] 154 | 155 | for i, t in enumerate(self.progress_bar(self.scheduler.timesteps)): 156 | # expand the latents if we are doing classifier free guidance 157 | latent_model_input = np.concatenate([latents] * 2) if do_classifier_free_guidance else latents 158 | latent_model_input = self.scheduler.scale_model_input(torch.from_numpy(latent_model_input), t) 159 | latent_model_input = latent_model_input.cpu().numpy() 160 | 161 | # predict the noise residual 162 | timestep = np.array([t], dtype=timestep_dtype) 163 | noise_pred = self.unet(sample=latent_model_input, timestep=timestep, encoder_hidden_states=prompt_embeds) 164 | noise_pred = noise_pred[0] 165 | 166 | # perform guidance 167 | if do_classifier_free_guidance: 168 | noise_pred_uncond, noise_pred_text = np.split(noise_pred, 2) 169 | noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) 170 | 171 | # compute the previous noisy sample x_t -> x_t-1 172 | scheduler_output = self.scheduler.step( 173 | torch.from_numpy(noise_pred), t, torch.from_numpy(latents), **extra_step_kwargs 174 | ) 175 | latents = scheduler_output.prev_sample.numpy() 176 | 177 | # call the callback, if provided 178 | 179 | if callback is not None and i % callback_steps == 0: 180 | cancel=callback(i, t, latents) 181 | if cancel: break 182 | 183 | latents = 1 / 0.18215 * latents 184 | #latents = 1 / 0.08333 * latents 185 | #print(type(latents)) 186 | #print(type(latents.shape)) 187 | #print(latents.shape) 188 | # image = self.vae_decoder(latent_sample=latents)[0] 189 | # it seems likes there is a strange result for using half-precision vae decoder if batchsize>1 190 | image = np.concatenate( 191 | [self.vae_decoder(latent_sample=latents[i : i + 1])[0] for i in range(latents.shape[0])] 192 | ) 193 | image = np.clip(image / 2 + 0.5, 0, 1) 194 | image = image.transpose((0, 2, 3, 1)) 195 | 196 | 197 | 198 | if self.safety_checker is not None: 199 | safety_checker_input = self.feature_extractor( 200 | self.numpy_to_pil(image), return_tensors="np" 201 | ).pixel_values.astype(image.dtype) 202 | 203 | images, has_nsfw_concept = [], [] 204 | for i in range(image.shape[0]): 205 | image_i, has_nsfw_concept_i = self.safety_checker( 206 | clip_input=safety_checker_input[i : i + 1], images=image[i : i + 1] 207 | ) 208 | images.append(image_i) 209 | has_nsfw_concept.append(has_nsfw_concept_i[0]) 210 | image = np.concatenate(images) 211 | else: 212 | has_nsfw_concept = None 213 | 214 | if output_type == "pil": 215 | image = self.numpy_to_pil(image) 216 | 217 | if not return_dict: 218 | return (image, has_nsfw_concept) 219 | 220 | #image.extend(images_saved) 221 | return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept) 222 | -------------------------------------------------------------------------------- /UI/ui_face_tools.py: -------------------------------------------------------------------------------- 1 | import gradio as gr 2 | import gc,os 3 | import numpy as np 4 | from Engine import pipelines_engines 5 | 6 | global facedect 7 | global faces 8 | global full_faces_info 9 | global new_faces 10 | global boxes 11 | global index 12 | global analyser 13 | analyser=None 14 | facedect = None 15 | faces = [] 16 | new_faces = [] 17 | boxes = [] 18 | index=0 19 | 20 | 21 | 22 | def show_input_image_area(): 23 | with gr.Row(): 24 | with gr.Column(variant="compact"): 25 | image_in = gr.Image(label="Input image", type="pil", elem_id="image_init") 26 | with gr.Column(variant="compact"): 27 | gallery_extracted_faces = gr.Gallery(label="Extracted Faces") 28 | facedect_btn = gr.Button("Extract faces information") 29 | select_face_btn = gr.Button("Select face") 30 | 31 | 32 | with gr.Row(): 33 | with gr.Column(variant="compact"): 34 | selected_face_in = gr.Image(label="Original Face", type="pil", elem_id="imagen2") 35 | with gr.Column(variant="compact"): 36 | with gr.Row(): 37 | face_in = gr.Image(label="Reference Face", type="pil",elem_id="imagen2") 38 | with gr.Row(): 39 | process_face_btn = gr.Button("Process Selected face",variant="primary") 40 | threshold = gr.Slider(1, 255, value=20, label="Threshold", info="Strength of pasted image(inversed), Increase until the pasted borders are not shown") 41 | 42 | 43 | 44 | with gr.Row(): 45 | selected_face_index = gr.Textbox(label="selected_face_index",value="", lines=1,visible=False, interactive=False) 46 | selected_face_gender = gr.Textbox(label="selected_face_gender",value="", lines=1,visible=False, interactive=True) 47 | selected_face_age = gr.Textbox(label="selected_face_age",value="", lines=1,visible=False, interactive=True) 48 | selected_face_box = gr.Textbox(label="selected_face_box",value="", lines=1,visible=False, interactive=False) 49 | with gr.Row(): 50 | new_face_info = gr.Textbox(label="New Face info",value="", lines=4,visible=False) 51 | with gr.Row(): 52 | update_face_btn = gr.Button("Update Gender & Age of Selected face",visible=False) 53 | with gr.Row(): 54 | new_face_out = gr.Image(label="Swapped Face& Create Mask",tool="sketch", type="pil", interactive=True,elem_id="imagen1") 55 | with gr.Row(): 56 | re_process_face_btn = gr.Button("Apply mask to processed face",variant="primary") 57 | with gr.Row(): 58 | new_face_out2 = gr.Image(label="Swapped Face", type="pil", interactive=True,elem_id="imagen1") 59 | with gr.Row(): 60 | paste_face_btn = gr.Button("Paste back new face",visible=False) 61 | unload_btn = gr.Button("Unload") 62 | 63 | facedect_btn.click(fn=face_detection, inputs=image_in, outputs=gallery_extracted_faces) 64 | select_face_btn.click(fn=select_face, inputs=None, outputs=[selected_face_in,selected_face_box,selected_face_index,selected_face_gender,selected_face_age]) 65 | process_face_btn.click(fn=swap_face, inputs=[image_in,selected_face_index,face_in,threshold], outputs=new_face_out) 66 | re_process_face_btn.click(fn=re_swap_face, inputs=[image_in,selected_face_index,face_in,threshold,new_face_out], outputs=new_face_out2) 67 | paste_face_btn.click(fn=paste_back_face, inputs=None) 68 | gallery_extracted_faces.select(fn=selected_index,inputs=None, outputs=None) 69 | unload_btn.click(fn=unload,inputs=None, outputs=None) 70 | update_face_btn.click(fn=update_face,inputs=[selected_face_index,selected_face_gender,selected_face_age], outputs=None) 71 | 72 | 73 | 74 | 75 | 76 | def update_face(index,gender,age): 77 | print("Update face") 78 | index=int(index) 79 | gender= int(gender) 80 | age= int(age) 81 | gender= 0 if gender==0 else 1 82 | global full_faces_info 83 | full_faces_info[index]['age']=age 84 | full_faces_info[index]['gender']=gender 85 | return 86 | 87 | def unload(): 88 | global analyser 89 | analyser=None 90 | 91 | 92 | def face_detection(image): 93 | print("face_detection") 94 | global faces 95 | global analyser 96 | global full_faces_info 97 | from Scripts.faceanalyser import faceanalyser 98 | if analyser==None: 99 | analyser=faceanalyser.face_analyser() 100 | 101 | full_faces_info=analyser.get(image) 102 | #print(full_faces_info[0].keys()) 103 | faces=[] 104 | image=np.asarray(image) 105 | from PIL import Image 106 | for face in full_faces_info: 107 | box=face['bbox'] 108 | #print(box) 109 | extracted_face = image[int(box[1]):int(box[3]), int(box[0]):int(box[2])] 110 | faces.append(Image.fromarray(extracted_face)) 111 | 112 | #image=Image.fromarray(image) 113 | #return image,faces 114 | return faces 115 | 116 | 117 | 118 | 119 | def select_face(): 120 | print("select_face") 121 | global faces 122 | global boxes 123 | global full_faces_info 124 | return faces[index],full_faces_info[index],index,full_faces_info[index]['gender'],full_faces_info[index]['age'] 125 | 126 | def re_swap_face(*args): #[image_in,selected_face_index,face_in,new_face_out] 127 | print("swap_face") 128 | import insightface 129 | global full_faces_info 130 | temp_frame=args[0] 131 | target_face=full_faces_info[int(args[1])] 132 | source_face=args[2] 133 | threshold=args[3] 134 | mask=args[4]["mask"].convert("RGB") 135 | 136 | global analyser 137 | source_face=analyser.get(source_face)[0] 138 | 139 | model_path = './Scripts/facedetector/inswapper_128_fp16.onnx' 140 | #face_processor = insightface.model_zoo.get_model(model_path, providers = ['DmlExecutionProvider']) 141 | face_processor = insightface.model_zoo.get_model(model_path, providers = ['CPUExecutionProvider']) 142 | #imagen=face_processor.get(np.asarray(temp_frame), target_face, source_face, paste_back = False) 143 | imagen=get_face(np.asarray(temp_frame), target_face, source_face,face_processor, paste_back = True,threshold=threshold,mask=np.asarray(mask),apply_mask=True) 144 | 145 | return imagen 146 | 147 | def swap_face(*args): #[image_in,selected_face_index,face_in] 148 | print("swap_face") 149 | import insightface 150 | global full_faces_info 151 | temp_frame=args[0] 152 | target_face=full_faces_info[int(args[1])] 153 | source_face=args[2] 154 | threshold=args[3] 155 | global analyser 156 | source_face=analyser.get(source_face)[0] 157 | 158 | model_path = './Scripts/facedetector/inswapper_128_fp16.onnx' 159 | #face_processor = insightface.model_zoo.get_model(model_path, providers = ['DmlExecutionProvider']) 160 | face_processor = insightface.model_zoo.get_model(model_path, providers = ['CPUExecutionProvider']) 161 | #imagen=face_processor.get(np.asarray(temp_frame), target_face, source_face, paste_back = False) 162 | imagen=get_face(np.asarray(temp_frame), target_face, source_face,face_processor, paste_back = True,threshold=threshold) 163 | 164 | return imagen 165 | 166 | def paste_back_face(): 167 | print("paste_back_face") 168 | return 169 | 170 | def selected_index(evt:gr.SelectData): 171 | #print("selected_index") 172 | global index 173 | index=evt.index 174 | return 175 | 176 | 177 | 178 | def get_face(img, target_face, source_face,face_processor, paste_back=True,threshold=20,mask=None,apply_mask=False): 179 | import cv2 180 | from insightface.utils import face_align 181 | 182 | #aimg, M = face_align.norm_crop2(img, target_face.kps, face_processor.input_size[0]) 183 | aimg_init, M = face_align.norm_crop2(img, target_face.kps, 672) 184 | aimg=cv2.resize(aimg_init,[128, 128], interpolation = cv2.INTER_AREA) 185 | blob = cv2.dnn.blobFromImage(aimg, 1.0 / face_processor.input_std, face_processor.input_size, 186 | (face_processor.input_mean, face_processor.input_mean, face_processor.input_mean), swapRB=True) 187 | latent = source_face.normed_embedding.reshape((1,-1)) 188 | latent = np.dot(latent, face_processor.emap) 189 | latent /= np.linalg.norm(latent) 190 | pred = face_processor.session.run(face_processor.output_names, {face_processor.input_names[0]: blob, face_processor.input_names[1]: latent})[0] 191 | img_fake = pred.transpose((0,2,3,1))[0] 192 | bgr_fake = np.clip(255 * img_fake, 0, 255).astype(np.uint8)[:,:,::-1] 193 | if not paste_back: 194 | return bgr_fake, M 195 | else: 196 | from Scripts import superresolution as SPR 197 | from PIL import Image 198 | aimg=aimg_init 199 | bgr_fake_PIL=SPR.superresolution_process_pil(Image.fromarray(bgr_fake)) 200 | bgr_fake= np.array(bgr_fake_PIL) 201 | 202 | target_img = img 203 | fake_diff = bgr_fake.astype(np.float32) - aimg.astype(np.float32) 204 | fake_diff = np.abs(fake_diff).mean(axis=2) 205 | fake_diff[:2,:] = 0 206 | fake_diff[-2:,:] = 0 207 | fake_diff[:,:2] = 0 208 | fake_diff[:,-2:] = 0 209 | IM = cv2.invertAffineTransform(M) 210 | img_white = np.full((aimg.shape[0],aimg.shape[1]), 255, dtype=np.float32) 211 | bgr_fake = cv2.warpAffine(bgr_fake, IM, (target_img.shape[1], target_img.shape[0]), borderValue=0.0) 212 | img_white = cv2.warpAffine(img_white, IM, (target_img.shape[1], target_img.shape[0]), borderValue=0.0) 213 | fake_diff = cv2.warpAffine(fake_diff, IM, (target_img.shape[1], target_img.shape[0]), borderValue=0.0) 214 | img_white[img_white>20] = 255 215 | #fthresh = 10 216 | fthresh = threshold 217 | 218 | fake_diff[fake_diff=fthresh] = 255 220 | img_mask = img_white 221 | mask_h_inds, mask_w_inds = np.where(img_mask==255) 222 | mask_h = np.max(mask_h_inds) - np.min(mask_h_inds) 223 | mask_w = np.max(mask_w_inds) - np.min(mask_w_inds) 224 | mask_size = int(np.sqrt(mask_h*mask_w)) 225 | k = max(mask_size//10, 10) 226 | #k = max(mask_size//20, 6) 227 | #k = 6 228 | cv2.imwrite("fake_diff.png", fake_diff) 229 | cv2.imwrite("img_mask.png", img_mask) 230 | if apply_mask: 231 | cv2.imwrite("mask.png", mask) 232 | 233 | #mask= np.asarray(mask.convert('RGB')) 234 | mask = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY) 235 | 236 | #img_mask=img_mask - mask 237 | fake_diff=fake_diff - mask 238 | fake_diff[fake_diff=fthresh] = 255 240 | 241 | cv2.imwrite("fake_diff.png", fake_diff) 242 | kernel = np.ones((k,k),np.uint8) 243 | img_mask = cv2.erode(img_mask,kernel,iterations = 1) 244 | kernel = np.ones((2,2),np.uint8) 245 | fake_diff = cv2.dilate(fake_diff,kernel,iterations = 1) 246 | k = max(mask_size//20, 5) 247 | #k = 3 248 | #k = 3 249 | kernel_size = (k, k) 250 | blur_size = tuple(2*i+1 for i in kernel_size) 251 | img_mask = cv2.GaussianBlur(img_mask, blur_size, 0) 252 | k = 5 253 | kernel_size = (k, k) 254 | blur_size = tuple(2*i+1 for i in kernel_size) 255 | fake_diff = cv2.GaussianBlur(fake_diff, blur_size, 0) 256 | img_mask /= 255 257 | fake_diff /= 255 258 | img_mask = fake_diff #uncommented from original file 259 | img_mask = np.reshape(img_mask, [img_mask.shape[0],img_mask.shape[1],1]) 260 | fake_merged = img_mask * bgr_fake + (1-img_mask) * target_img.astype(np.float32) 261 | fake_merged = fake_merged.astype(np.uint8) 262 | return fake_merged 263 | 264 | 265 | --------------------------------------------------------------------------------