use the new callback for script unloaded to stop the script from having effect after it's unloaded

master
AUTOMATIC 2023-01-06 13:59:30 +03:00
parent 411889ca60
commit 2624e5dd49
1 changed files with 8 additions and 1 deletions

View File

@ -1,5 +1,5 @@
from modules import scripts, script_callbacks
from modules import scripts, script_callbacks, shared
import aesthetic_clip
import gradio as gr
@ -41,6 +41,12 @@ def on_model_loaded(sd_model):
sd_model.cond_stage_model.process_tokens = aesthetic
def on_script_unloaded():
cond_stage_model = shared.sd_model.cond_stage_model
if type(cond_stage_model.process_tokens) == aesthetic_clip.AestheticCLIP:
cond_stage_model.process_tokens = cond_stage_model.process_tokens.process_tokens
def on_ui_tabs():
with gr.Blocks(analytics_enabled=False) as aesthetic_interface:
with gr.Row().style(equal_height=False):
@ -89,5 +95,6 @@ def on_ui_tabs():
return [(aesthetic_interface, "Create aesthetic embedding", "aesthetic_interface")]
script_callbacks.on_script_unloaded(on_script_unloaded)
script_callbacks.on_model_loaded(on_model_loaded)
script_callbacks.on_ui_tabs(on_ui_tabs)