cymic/Waifu_Diffusion_Webui

1import os2import threading3import time4import importlib5import signal6import threading7 8from modules.paths import script_path9 10from modules import devices, sd_samplers11import modules.codeformer_model as codeformer12import modules.extras13import modules.face_restoration14import modules.gfpgan_model as gfpgan15import modules.img2img16 17import modules.lowvram18import modules.paths19import modules.scripts20import modules.sd_hijack21import modules.sd_models22import modules.shared as shared23import modules.txt2img24 25import modules.ui26from modules import devices27from modules import modelloader28from modules.paths import script_path29from modules.shared import cmd_opts30 31modelloader.cleanup_models()32modules.sd_models.setup_model()33codeformer.setup_model(cmd_opts.codeformer_models_path)34gfpgan.setup_model(cmd_opts.gfpgan_models_path)35shared.face_restorers.append(modules.face_restoration.FaceRestoration())36modelloader.load_upscalers()37queue_lock = threading.Lock()38 39 40def wrap_queued_call(func):41    def f(*args, **kwargs):42        with queue_lock:43            res = func(*args, **kwargs)44 45        return res46 47    return f48 49 50def wrap_gradio_gpu_call(func, extra_outputs=None):51    def f(*args, **kwargs):52        devices.torch_gc()53 54        shared.state.sampling_step = 055        shared.state.job_count = -156        shared.state.job_no = 057        shared.state.job_timestamp = shared.state.get_job_timestamp()58        shared.state.current_latent = None59        shared.state.current_image = None60        shared.state.current_image_sampling_step = 061        shared.state.interrupted = False62        shared.state.textinfo = None63 64        with queue_lock:65            res = func(*args, **kwargs)66 67        shared.state.job = ""68        shared.state.job_count = 069 70        devices.torch_gc()71 72        return res73 74    return modules.ui.wrap_gradio_call(f, extra_outputs=extra_outputs)75 76 77modules.scripts.load_scripts(os.path.join(script_path, "scripts"))78 79shared.sd_model = modules.sd_models.load_model()80shared.opts.onchange("sd_model_checkpoint", wrap_queued_call(lambda: modules.sd_models.reload_model_weights(shared.sd_model)))81 82 83def webui():84    # make the program just exit at ctrl+c without waiting for anything85    def sigint_handler(sig, frame):86        print(f'Interrupted with signal {sig} in {frame}')87        os._exit(0)88 89    signal.signal(signal.SIGINT, sigint_handler)90 91    while 1:92 93        demo = modules.ui.create_ui(wrap_gradio_gpu_call=wrap_gradio_gpu_call)94        95        demo.launch(96            share=cmd_opts.share,97            server_name="0.0.0.0" if cmd_opts.listen else None,98            server_port=cmd_opts.port,99            debug=cmd_opts.gradio_debug,100            auth=[tuple(cred.split(':')) for cred in cmd_opts.gradio_auth.strip('"').split(',')] if cmd_opts.gradio_auth else None,101            inbrowser=cmd_opts.autolaunch,102            prevent_thread_lock=True103        )104 105        while 1:106            time.sleep(0.5)107            if getattr(demo, 'do_restart', False):108                time.sleep(0.5)109                demo.close()110                time.sleep(0.5)111                break112 113        sd_samplers.set_samplers()114 115        print('Reloading Custom Scripts')116        modules.scripts.reload_scripts(os.path.join(script_path, "scripts"))117        print('Reloading modules: modules.ui')118        importlib.reload(modules.ui)119        print('Restarting Gradio')120 121 122 123if __name__ == "__main__":124    webui()125