File size: 3,912 Bytes
0163a2c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 |
from modules import scripts, script_callbacks, shared
import aesthetic_clip
import gradio as gr
aesthetic = aesthetic_clip.AestheticCLIP()
aesthetic_imgs_components = []
class AestheticScript(scripts.Script):
def title(self):
return "Aesthetic embeddings"
def show(self, is_img2img):
return scripts.AlwaysVisible
def ui(self, is_img2img):
aesthetic_weight, aesthetic_steps, aesthetic_lr, aesthetic_slerp, aesthetic_imgs, aesthetic_imgs_text, aesthetic_slerp_angle, aesthetic_text_negative = aesthetic_clip.create_ui()
self.infotext_fields = [
(aesthetic_lr, "Aesthetic LR"),
(aesthetic_weight, "Aesthetic weight"),
(aesthetic_steps, "Aesthetic steps"),
(aesthetic_imgs, "Aesthetic embedding"),
(aesthetic_slerp, "Aesthetic slerp"),
(aesthetic_imgs_text, "Aesthetic text"),
(aesthetic_text_negative, "Aesthetic text negative"),
(aesthetic_slerp_angle, "Aesthetic slerp angle"),
]
aesthetic_imgs_components.append(aesthetic_imgs)
return [aesthetic_weight, aesthetic_steps, aesthetic_lr, aesthetic_slerp, aesthetic_imgs, aesthetic_imgs_text, aesthetic_slerp_angle, aesthetic_text_negative]
def process(self, p, aesthetic_weight, aesthetic_steps, aesthetic_lr, aesthetic_slerp, aesthetic_imgs, aesthetic_imgs_text, aesthetic_slerp_angle, aesthetic_text_negative):
aesthetic.set_aesthetic_params(p, float(aesthetic_lr), float(aesthetic_weight), int(aesthetic_steps), aesthetic_imgs, aesthetic_slerp, aesthetic_imgs_text, aesthetic_slerp_angle, aesthetic_text_negative)
def on_model_loaded(sd_model):
aesthetic.process_tokens = sd_model.cond_stage_model.process_tokens
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):
with gr.Column(variant='panel'):
gr.HTML(value="Create an aesthetic embedding out of any number of images")
new_embedding_name_ae = gr.Textbox(label="Name")
process_src_ae = gr.Textbox(label='Source directory')
batch_ae = gr.Slider(minimum=1, maximum=1024, step=1, label="Batch size", value=256)
with gr.Row():
with gr.Column(scale=3):
gr.HTML(value="")
with gr.Column():
create_embedding_ae = gr.Button(value="Create images embedding", variant='primary')
with gr.Column():
output = gr.Text(value="", show_label=False)
dropdown_components = aesthetic_imgs_components.copy()
def generate_embs(*args):
res = aesthetic_clip.generate_imgs_embd(*args)
aesthetic_clip.update_aesthetic_embeddings()
updates = [gr.Dropdown.update(choices=sorted(aesthetic_clip.aesthetic_embeddings.keys())) for _ in range(len(dropdown_components))]
return [*updates, res]
create_embedding_ae.click(
fn=generate_embs,
inputs=[
new_embedding_name_ae,
process_src_ae,
batch_ae
],
outputs=[
*dropdown_components,
output
]
)
aesthetic_imgs_components.clear()
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)
|