e82ea20299
Child classes only get populated to __subclassess__ when they are imported. We don't actually need to import any of them to webui any more, so clean up webUI imports and make sure loader imports children. Also, fix command line paths not actually being passed to the scalers.
102 lines
2.9 KiB
Python
102 lines
2.9 KiB
Python
import os
|
|
import signal
|
|
import threading
|
|
|
|
import modules.codeformer_model as codeformer
|
|
import modules.extras
|
|
import modules.face_restoration
|
|
import modules.gfpgan_model as gfpgan
|
|
import modules.img2img
|
|
import modules.lowvram
|
|
import modules.paths
|
|
import modules.scripts
|
|
import modules.sd_hijack
|
|
import modules.sd_models
|
|
import modules.shared as shared
|
|
import modules.txt2img
|
|
import modules.ui
|
|
from modules import devices
|
|
from modules import modelloader
|
|
from modules.paths import script_path
|
|
from modules.shared import cmd_opts
|
|
|
|
modelloader.cleanup_models()
|
|
modules.sd_models.setup_model(cmd_opts.ckpt_dir)
|
|
codeformer.setup_model(cmd_opts.codeformer_models_path)
|
|
gfpgan.setup_model(cmd_opts.gfpgan_models_path)
|
|
shared.face_restorers.append(modules.face_restoration.FaceRestoration())
|
|
modelloader.load_upscalers()
|
|
queue_lock = threading.Lock()
|
|
|
|
|
|
def wrap_queued_call(func):
|
|
def f(*args, **kwargs):
|
|
with queue_lock:
|
|
res = func(*args, **kwargs)
|
|
|
|
return res
|
|
|
|
return f
|
|
|
|
|
|
def wrap_gradio_gpu_call(func):
|
|
def f(*args, **kwargs):
|
|
devices.torch_gc()
|
|
|
|
shared.state.sampling_step = 0
|
|
shared.state.job_count = -1
|
|
shared.state.job_no = 0
|
|
shared.state.job_timestamp = shared.state.get_job_timestamp()
|
|
shared.state.current_latent = None
|
|
shared.state.current_image = None
|
|
shared.state.current_image_sampling_step = 0
|
|
shared.state.interrupted = False
|
|
|
|
with queue_lock:
|
|
res = func(*args, **kwargs)
|
|
|
|
shared.state.job = ""
|
|
shared.state.job_count = 0
|
|
|
|
devices.torch_gc()
|
|
|
|
return res
|
|
|
|
return modules.ui.wrap_gradio_call(f)
|
|
|
|
|
|
modules.scripts.load_scripts(os.path.join(script_path, "scripts"))
|
|
|
|
shared.sd_model = modules.sd_models.load_model()
|
|
shared.opts.onchange("sd_model_checkpoint", wrap_queued_call(lambda: modules.sd_models.reload_model_weights(shared.sd_model)))
|
|
|
|
|
|
def webui():
|
|
# make the program just exit at ctrl+c without waiting for anything
|
|
def sigint_handler(sig, frame):
|
|
print(f'Interrupted with signal {sig} in {frame}')
|
|
os._exit(0)
|
|
|
|
signal.signal(signal.SIGINT, sigint_handler)
|
|
|
|
demo = modules.ui.create_ui(
|
|
txt2img=wrap_gradio_gpu_call(modules.txt2img.txt2img),
|
|
img2img=wrap_gradio_gpu_call(modules.img2img.img2img),
|
|
run_extras=wrap_gradio_gpu_call(modules.extras.run_extras),
|
|
run_pnginfo=modules.extras.run_pnginfo,
|
|
run_modelmerger=modules.extras.run_modelmerger
|
|
)
|
|
|
|
demo.launch(
|
|
share=cmd_opts.share,
|
|
server_name="0.0.0.0" if cmd_opts.listen else None,
|
|
server_port=cmd_opts.port,
|
|
debug=cmd_opts.gradio_debug,
|
|
auth=[tuple(cred.split(':')) for cred in cmd_opts.gradio_auth.strip('"').split(',')] if cmd_opts.gradio_auth else None,
|
|
inbrowser=cmd_opts.autolaunch,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
webui()
|