removed aesthetic gradients as built-in
added support for extensions
This commit is contained in:
parent
26d1073745
commit
2b91251637
14 changed files with 249 additions and 410 deletions
0
extensions/put extension here.txt
Normal file
0
extensions/put extension here.txt
Normal file
|
@ -1,241 +0,0 @@
|
||||||
import copy
|
|
||||||
import itertools
|
|
||||||
import os
|
|
||||||
from pathlib import Path
|
|
||||||
import html
|
|
||||||
import gc
|
|
||||||
|
|
||||||
import gradio as gr
|
|
||||||
import torch
|
|
||||||
from PIL import Image
|
|
||||||
from torch import optim
|
|
||||||
|
|
||||||
from modules import shared
|
|
||||||
from transformers import CLIPModel, CLIPProcessor, CLIPTokenizer
|
|
||||||
from tqdm.auto import tqdm, trange
|
|
||||||
from modules.shared import opts, device
|
|
||||||
|
|
||||||
|
|
||||||
def get_all_images_in_folder(folder):
|
|
||||||
return [os.path.join(folder, f) for f in os.listdir(folder) if
|
|
||||||
os.path.isfile(os.path.join(folder, f)) and check_is_valid_image_file(f)]
|
|
||||||
|
|
||||||
|
|
||||||
def check_is_valid_image_file(filename):
|
|
||||||
return filename.lower().endswith(('.png', '.jpg', '.jpeg', ".gif", ".tiff", ".webp"))
|
|
||||||
|
|
||||||
|
|
||||||
def batched(dataset, total, n=1):
|
|
||||||
for ndx in range(0, total, n):
|
|
||||||
yield [dataset.__getitem__(i) for i in range(ndx, min(ndx + n, total))]
|
|
||||||
|
|
||||||
|
|
||||||
def iter_to_batched(iterable, n=1):
|
|
||||||
it = iter(iterable)
|
|
||||||
while True:
|
|
||||||
chunk = tuple(itertools.islice(it, n))
|
|
||||||
if not chunk:
|
|
||||||
return
|
|
||||||
yield chunk
|
|
||||||
|
|
||||||
|
|
||||||
def create_ui():
|
|
||||||
import modules.ui
|
|
||||||
|
|
||||||
with gr.Group():
|
|
||||||
with gr.Accordion("Open for Clip Aesthetic!", open=False):
|
|
||||||
with gr.Row():
|
|
||||||
aesthetic_weight = gr.Slider(minimum=0, maximum=1, step=0.01, label="Aesthetic weight",
|
|
||||||
value=0.9)
|
|
||||||
aesthetic_steps = gr.Slider(minimum=0, maximum=50, step=1, label="Aesthetic steps", value=5)
|
|
||||||
|
|
||||||
with gr.Row():
|
|
||||||
aesthetic_lr = gr.Textbox(label='Aesthetic learning rate',
|
|
||||||
placeholder="Aesthetic learning rate", value="0.0001")
|
|
||||||
aesthetic_slerp = gr.Checkbox(label="Slerp interpolation", value=False)
|
|
||||||
aesthetic_imgs = gr.Dropdown(sorted(shared.aesthetic_embeddings.keys()),
|
|
||||||
label="Aesthetic imgs embedding",
|
|
||||||
value="None")
|
|
||||||
|
|
||||||
modules.ui.create_refresh_button(aesthetic_imgs, shared.update_aesthetic_embeddings, lambda: {"choices": sorted(shared.aesthetic_embeddings.keys())}, "refresh_aesthetic_embeddings")
|
|
||||||
|
|
||||||
with gr.Row():
|
|
||||||
aesthetic_imgs_text = gr.Textbox(label='Aesthetic text for imgs',
|
|
||||||
placeholder="This text is used to rotate the feature space of the imgs embs",
|
|
||||||
value="")
|
|
||||||
aesthetic_slerp_angle = gr.Slider(label='Slerp angle', minimum=0, maximum=1, step=0.01,
|
|
||||||
value=0.1)
|
|
||||||
aesthetic_text_negative = gr.Checkbox(label="Is negative text", value=False)
|
|
||||||
|
|
||||||
return aesthetic_weight, aesthetic_steps, aesthetic_lr, aesthetic_slerp, aesthetic_imgs, aesthetic_imgs_text, aesthetic_slerp_angle, aesthetic_text_negative
|
|
||||||
|
|
||||||
|
|
||||||
aesthetic_clip_model = None
|
|
||||||
|
|
||||||
|
|
||||||
def aesthetic_clip():
|
|
||||||
global aesthetic_clip_model
|
|
||||||
|
|
||||||
if aesthetic_clip_model is None or aesthetic_clip_model.name_or_path != shared.sd_model.cond_stage_model.wrapped.transformer.name_or_path:
|
|
||||||
aesthetic_clip_model = CLIPModel.from_pretrained(shared.sd_model.cond_stage_model.wrapped.transformer.name_or_path)
|
|
||||||
aesthetic_clip_model.cpu()
|
|
||||||
|
|
||||||
return aesthetic_clip_model
|
|
||||||
|
|
||||||
|
|
||||||
def generate_imgs_embd(name, folder, batch_size):
|
|
||||||
model = aesthetic_clip().to(device)
|
|
||||||
processor = CLIPProcessor.from_pretrained(model.name_or_path)
|
|
||||||
|
|
||||||
with torch.no_grad():
|
|
||||||
embs = []
|
|
||||||
for paths in tqdm(iter_to_batched(get_all_images_in_folder(folder), batch_size),
|
|
||||||
desc=f"Generating embeddings for {name}"):
|
|
||||||
if shared.state.interrupted:
|
|
||||||
break
|
|
||||||
inputs = processor(images=[Image.open(path) for path in paths], return_tensors="pt").to(device)
|
|
||||||
outputs = model.get_image_features(**inputs).cpu()
|
|
||||||
embs.append(torch.clone(outputs))
|
|
||||||
inputs.to("cpu")
|
|
||||||
del inputs, outputs
|
|
||||||
|
|
||||||
embs = torch.cat(embs, dim=0).mean(dim=0, keepdim=True)
|
|
||||||
|
|
||||||
# The generated embedding will be located here
|
|
||||||
path = str(Path(shared.cmd_opts.aesthetic_embeddings_dir) / f"{name}.pt")
|
|
||||||
torch.save(embs, path)
|
|
||||||
|
|
||||||
model.cpu()
|
|
||||||
del processor
|
|
||||||
del embs
|
|
||||||
gc.collect()
|
|
||||||
torch.cuda.empty_cache()
|
|
||||||
res = f"""
|
|
||||||
Done generating embedding for {name}!
|
|
||||||
Aesthetic embedding saved to {html.escape(path)}
|
|
||||||
"""
|
|
||||||
shared.update_aesthetic_embeddings()
|
|
||||||
return gr.Dropdown.update(choices=sorted(shared.aesthetic_embeddings.keys()), label="Imgs embedding",
|
|
||||||
value="None"), \
|
|
||||||
gr.Dropdown.update(choices=sorted(shared.aesthetic_embeddings.keys()),
|
|
||||||
label="Imgs embedding",
|
|
||||||
value="None"), res, ""
|
|
||||||
|
|
||||||
|
|
||||||
def slerp(low, high, val):
|
|
||||||
low_norm = low / torch.norm(low, dim=1, keepdim=True)
|
|
||||||
high_norm = high / torch.norm(high, dim=1, keepdim=True)
|
|
||||||
omega = torch.acos((low_norm * high_norm).sum(1))
|
|
||||||
so = torch.sin(omega)
|
|
||||||
res = (torch.sin((1.0 - val) * omega) / so).unsqueeze(1) * low + (torch.sin(val * omega) / so).unsqueeze(1) * high
|
|
||||||
return res
|
|
||||||
|
|
||||||
|
|
||||||
class AestheticCLIP:
|
|
||||||
def __init__(self):
|
|
||||||
self.skip = False
|
|
||||||
self.aesthetic_steps = 0
|
|
||||||
self.aesthetic_weight = 0
|
|
||||||
self.aesthetic_lr = 0
|
|
||||||
self.slerp = False
|
|
||||||
self.aesthetic_text_negative = ""
|
|
||||||
self.aesthetic_slerp_angle = 0
|
|
||||||
self.aesthetic_imgs_text = ""
|
|
||||||
|
|
||||||
self.image_embs_name = None
|
|
||||||
self.image_embs = None
|
|
||||||
self.load_image_embs(None)
|
|
||||||
|
|
||||||
def set_aesthetic_params(self, p, aesthetic_lr=0, aesthetic_weight=0, aesthetic_steps=0, image_embs_name=None,
|
|
||||||
aesthetic_slerp=True, aesthetic_imgs_text="",
|
|
||||||
aesthetic_slerp_angle=0.15,
|
|
||||||
aesthetic_text_negative=False):
|
|
||||||
self.aesthetic_imgs_text = aesthetic_imgs_text
|
|
||||||
self.aesthetic_slerp_angle = aesthetic_slerp_angle
|
|
||||||
self.aesthetic_text_negative = aesthetic_text_negative
|
|
||||||
self.slerp = aesthetic_slerp
|
|
||||||
self.aesthetic_lr = aesthetic_lr
|
|
||||||
self.aesthetic_weight = aesthetic_weight
|
|
||||||
self.aesthetic_steps = aesthetic_steps
|
|
||||||
self.load_image_embs(image_embs_name)
|
|
||||||
|
|
||||||
if self.image_embs_name is not None:
|
|
||||||
p.extra_generation_params.update({
|
|
||||||
"Aesthetic LR": aesthetic_lr,
|
|
||||||
"Aesthetic weight": aesthetic_weight,
|
|
||||||
"Aesthetic steps": aesthetic_steps,
|
|
||||||
"Aesthetic embedding": self.image_embs_name,
|
|
||||||
"Aesthetic slerp": aesthetic_slerp,
|
|
||||||
"Aesthetic text": aesthetic_imgs_text,
|
|
||||||
"Aesthetic text negative": aesthetic_text_negative,
|
|
||||||
"Aesthetic slerp angle": aesthetic_slerp_angle,
|
|
||||||
})
|
|
||||||
|
|
||||||
def set_skip(self, skip):
|
|
||||||
self.skip = skip
|
|
||||||
|
|
||||||
def load_image_embs(self, image_embs_name):
|
|
||||||
if image_embs_name is None or len(image_embs_name) == 0 or image_embs_name == "None":
|
|
||||||
image_embs_name = None
|
|
||||||
self.image_embs_name = None
|
|
||||||
if image_embs_name is not None and self.image_embs_name != image_embs_name:
|
|
||||||
self.image_embs_name = image_embs_name
|
|
||||||
self.image_embs = torch.load(shared.aesthetic_embeddings[self.image_embs_name], map_location=device)
|
|
||||||
self.image_embs /= self.image_embs.norm(dim=-1, keepdim=True)
|
|
||||||
self.image_embs.requires_grad_(False)
|
|
||||||
|
|
||||||
def __call__(self, z, remade_batch_tokens):
|
|
||||||
if not self.skip and self.aesthetic_steps != 0 and self.aesthetic_lr != 0 and self.aesthetic_weight != 0 and self.image_embs_name is not None:
|
|
||||||
tokenizer = shared.sd_model.cond_stage_model.tokenizer
|
|
||||||
if not opts.use_old_emphasis_implementation:
|
|
||||||
remade_batch_tokens = [
|
|
||||||
[tokenizer.bos_token_id] + x[:75] + [tokenizer.eos_token_id] for x in
|
|
||||||
remade_batch_tokens]
|
|
||||||
|
|
||||||
tokens = torch.asarray(remade_batch_tokens).to(device)
|
|
||||||
|
|
||||||
model = copy.deepcopy(aesthetic_clip()).to(device)
|
|
||||||
model.requires_grad_(True)
|
|
||||||
if self.aesthetic_imgs_text is not None and len(self.aesthetic_imgs_text) > 0:
|
|
||||||
text_embs_2 = model.get_text_features(
|
|
||||||
**tokenizer([self.aesthetic_imgs_text], padding=True, return_tensors="pt").to(device))
|
|
||||||
if self.aesthetic_text_negative:
|
|
||||||
text_embs_2 = self.image_embs - text_embs_2
|
|
||||||
text_embs_2 /= text_embs_2.norm(dim=-1, keepdim=True)
|
|
||||||
img_embs = slerp(self.image_embs, text_embs_2, self.aesthetic_slerp_angle)
|
|
||||||
else:
|
|
||||||
img_embs = self.image_embs
|
|
||||||
|
|
||||||
with torch.enable_grad():
|
|
||||||
|
|
||||||
# We optimize the model to maximize the similarity
|
|
||||||
optimizer = optim.Adam(
|
|
||||||
model.text_model.parameters(), lr=self.aesthetic_lr
|
|
||||||
)
|
|
||||||
|
|
||||||
for _ in trange(self.aesthetic_steps, desc="Aesthetic optimization"):
|
|
||||||
text_embs = model.get_text_features(input_ids=tokens)
|
|
||||||
text_embs = text_embs / text_embs.norm(dim=-1, keepdim=True)
|
|
||||||
sim = text_embs @ img_embs.T
|
|
||||||
loss = -sim
|
|
||||||
optimizer.zero_grad()
|
|
||||||
loss.mean().backward()
|
|
||||||
optimizer.step()
|
|
||||||
|
|
||||||
zn = model.text_model(input_ids=tokens, output_hidden_states=-opts.CLIP_stop_at_last_layers)
|
|
||||||
if opts.CLIP_stop_at_last_layers > 1:
|
|
||||||
zn = zn.hidden_states[-opts.CLIP_stop_at_last_layers]
|
|
||||||
zn = model.text_model.final_layer_norm(zn)
|
|
||||||
else:
|
|
||||||
zn = zn.last_hidden_state
|
|
||||||
model.cpu()
|
|
||||||
del model
|
|
||||||
gc.collect()
|
|
||||||
torch.cuda.empty_cache()
|
|
||||||
zn = torch.concat([zn[77 * i:77 * (i + 1)] for i in range(max(z.shape[1] // 77, 1))], 1)
|
|
||||||
if self.slerp:
|
|
||||||
z = slerp(z, zn, self.aesthetic_weight)
|
|
||||||
else:
|
|
||||||
z = z * (1 - self.aesthetic_weight) + zn * self.aesthetic_weight
|
|
||||||
|
|
||||||
return z
|
|
|
@ -56,7 +56,7 @@ def process_batch(p, input_dir, output_dir, args):
|
||||||
processed_image.save(os.path.join(output_dir, filename))
|
processed_image.save(os.path.join(output_dir, filename))
|
||||||
|
|
||||||
|
|
||||||
def img2img(mode: int, prompt: str, negative_prompt: str, prompt_style: str, prompt_style2: str, init_img, init_img_with_mask, init_img_inpaint, init_mask_inpaint, mask_mode, steps: int, sampler_index: int, mask_blur: int, inpainting_fill: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, denoising_strength: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, seed_enable_extras: bool, height: int, width: int, resize_mode: int, inpaint_full_res: bool, inpaint_full_res_padding: int, inpainting_mask_invert: int, img2img_batch_input_dir: str, img2img_batch_output_dir: str, aesthetic_lr=0, aesthetic_weight=0, aesthetic_steps=0, aesthetic_imgs=None, aesthetic_slerp=False, aesthetic_imgs_text="", aesthetic_slerp_angle=0.15, aesthetic_text_negative=False, *args):
|
def img2img(mode: int, prompt: str, negative_prompt: str, prompt_style: str, prompt_style2: str, init_img, init_img_with_mask, init_img_inpaint, init_mask_inpaint, mask_mode, steps: int, sampler_index: int, mask_blur: int, inpainting_fill: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, denoising_strength: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, seed_enable_extras: bool, height: int, width: int, resize_mode: int, inpaint_full_res: bool, inpaint_full_res_padding: int, inpainting_mask_invert: int, img2img_batch_input_dir: str, img2img_batch_output_dir: str, *args):
|
||||||
is_inpaint = mode == 1
|
is_inpaint = mode == 1
|
||||||
is_batch = mode == 2
|
is_batch = mode == 2
|
||||||
|
|
||||||
|
@ -109,7 +109,8 @@ def img2img(mode: int, prompt: str, negative_prompt: str, prompt_style: str, pro
|
||||||
inpainting_mask_invert=inpainting_mask_invert,
|
inpainting_mask_invert=inpainting_mask_invert,
|
||||||
)
|
)
|
||||||
|
|
||||||
shared.aesthetic_clip.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)
|
p.scripts = modules.scripts.scripts_txt2img
|
||||||
|
p.script_args = args
|
||||||
|
|
||||||
if shared.cmd_opts.enable_console_prompts:
|
if shared.cmd_opts.enable_console_prompts:
|
||||||
print(f"\nimg2img: {prompt}", file=shared.progress_print_out)
|
print(f"\nimg2img: {prompt}", file=shared.progress_print_out)
|
||||||
|
|
|
@ -104,6 +104,12 @@ class StableDiffusionProcessing():
|
||||||
self.seed_resize_from_h = 0
|
self.seed_resize_from_h = 0
|
||||||
self.seed_resize_from_w = 0
|
self.seed_resize_from_w = 0
|
||||||
|
|
||||||
|
self.scripts = None
|
||||||
|
self.script_args = None
|
||||||
|
self.all_prompts = None
|
||||||
|
self.all_seeds = None
|
||||||
|
self.all_subseeds = None
|
||||||
|
|
||||||
|
|
||||||
def init(self, all_prompts, all_seeds, all_subseeds):
|
def init(self, all_prompts, all_seeds, all_subseeds):
|
||||||
pass
|
pass
|
||||||
|
@ -350,32 +356,35 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
|
||||||
shared.prompt_styles.apply_styles(p)
|
shared.prompt_styles.apply_styles(p)
|
||||||
|
|
||||||
if type(p.prompt) == list:
|
if type(p.prompt) == list:
|
||||||
all_prompts = p.prompt
|
p.all_prompts = p.prompt
|
||||||
else:
|
else:
|
||||||
all_prompts = p.batch_size * p.n_iter * [p.prompt]
|
p.all_prompts = p.batch_size * p.n_iter * [p.prompt]
|
||||||
|
|
||||||
if type(seed) == list:
|
if type(seed) == list:
|
||||||
all_seeds = seed
|
p.all_seeds = seed
|
||||||
else:
|
else:
|
||||||
all_seeds = [int(seed) + (x if p.subseed_strength == 0 else 0) for x in range(len(all_prompts))]
|
p.all_seeds = [int(seed) + (x if p.subseed_strength == 0 else 0) for x in range(len(p.all_prompts))]
|
||||||
|
|
||||||
if type(subseed) == list:
|
if type(subseed) == list:
|
||||||
all_subseeds = subseed
|
p.all_subseeds = subseed
|
||||||
else:
|
else:
|
||||||
all_subseeds = [int(subseed) + x for x in range(len(all_prompts))]
|
p.all_subseeds = [int(subseed) + x for x in range(len(p.all_prompts))]
|
||||||
|
|
||||||
def infotext(iteration=0, position_in_batch=0):
|
def infotext(iteration=0, position_in_batch=0):
|
||||||
return create_infotext(p, all_prompts, all_seeds, all_subseeds, comments, iteration, position_in_batch)
|
return create_infotext(p, p.all_prompts, p.all_seeds, p.all_subseeds, comments, iteration, position_in_batch)
|
||||||
|
|
||||||
if os.path.exists(cmd_opts.embeddings_dir) and not p.do_not_reload_embeddings:
|
if os.path.exists(cmd_opts.embeddings_dir) and not p.do_not_reload_embeddings:
|
||||||
model_hijack.embedding_db.load_textual_inversion_embeddings()
|
model_hijack.embedding_db.load_textual_inversion_embeddings()
|
||||||
|
|
||||||
|
if p.scripts is not None:
|
||||||
|
p.scripts.run_alwayson_scripts(p)
|
||||||
|
|
||||||
infotexts = []
|
infotexts = []
|
||||||
output_images = []
|
output_images = []
|
||||||
|
|
||||||
with torch.no_grad(), p.sd_model.ema_scope():
|
with torch.no_grad(), p.sd_model.ema_scope():
|
||||||
with devices.autocast():
|
with devices.autocast():
|
||||||
p.init(all_prompts, all_seeds, all_subseeds)
|
p.init(p.all_prompts, p.all_seeds, p.all_subseeds)
|
||||||
|
|
||||||
if state.job_count == -1:
|
if state.job_count == -1:
|
||||||
state.job_count = p.n_iter
|
state.job_count = p.n_iter
|
||||||
|
@ -387,9 +396,9 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
|
||||||
if state.interrupted:
|
if state.interrupted:
|
||||||
break
|
break
|
||||||
|
|
||||||
prompts = all_prompts[n * p.batch_size:(n + 1) * p.batch_size]
|
prompts = p.all_prompts[n * p.batch_size:(n + 1) * p.batch_size]
|
||||||
seeds = all_seeds[n * p.batch_size:(n + 1) * p.batch_size]
|
seeds = p.all_seeds[n * p.batch_size:(n + 1) * p.batch_size]
|
||||||
subseeds = all_subseeds[n * p.batch_size:(n + 1) * p.batch_size]
|
subseeds = p.all_subseeds[n * p.batch_size:(n + 1) * p.batch_size]
|
||||||
|
|
||||||
if (len(prompts) == 0):
|
if (len(prompts) == 0):
|
||||||
break
|
break
|
||||||
|
@ -490,10 +499,10 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
|
||||||
index_of_first_image = 1
|
index_of_first_image = 1
|
||||||
|
|
||||||
if opts.grid_save:
|
if opts.grid_save:
|
||||||
images.save_image(grid, p.outpath_grids, "grid", all_seeds[0], all_prompts[0], opts.grid_format, info=infotext(), short_filename=not opts.grid_extended_filename, p=p, grid=True)
|
images.save_image(grid, p.outpath_grids, "grid", p.all_seeds[0], p.all_prompts[0], opts.grid_format, info=infotext(), short_filename=not opts.grid_extended_filename, p=p, grid=True)
|
||||||
|
|
||||||
devices.torch_gc()
|
devices.torch_gc()
|
||||||
return Processed(p, output_images, all_seeds[0], infotext() + "".join(["\n\n" + x for x in comments]), subseed=all_subseeds[0], all_prompts=all_prompts, all_seeds=all_seeds, all_subseeds=all_subseeds, index_of_first_image=index_of_first_image, infotexts=infotexts)
|
return Processed(p, output_images, p.all_seeds[0], infotext() + "".join(["\n\n" + x for x in comments]), subseed=p.all_subseeds[0], all_prompts=p.all_prompts, all_seeds=p.all_seeds, all_subseeds=p.all_subseeds, index_of_first_image=index_of_first_image, infotexts=infotexts)
|
||||||
|
|
||||||
|
|
||||||
class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
|
class StableDiffusionProcessingTxt2Img(StableDiffusionProcessing):
|
||||||
|
|
42
modules/script_callbacks.py
Normal file
42
modules/script_callbacks.py
Normal file
|
@ -0,0 +1,42 @@
|
||||||
|
|
||||||
|
callbacks_model_loaded = []
|
||||||
|
callbacks_ui_tabs = []
|
||||||
|
|
||||||
|
|
||||||
|
def clear_callbacks():
|
||||||
|
callbacks_model_loaded.clear()
|
||||||
|
callbacks_ui_tabs.clear()
|
||||||
|
|
||||||
|
|
||||||
|
def model_loaded_callback(sd_model):
|
||||||
|
for callback in callbacks_model_loaded:
|
||||||
|
callback(sd_model)
|
||||||
|
|
||||||
|
|
||||||
|
def ui_tabs_callback():
|
||||||
|
res = []
|
||||||
|
|
||||||
|
for callback in callbacks_ui_tabs:
|
||||||
|
res += callback() or []
|
||||||
|
|
||||||
|
return res
|
||||||
|
|
||||||
|
|
||||||
|
def on_model_loaded(callback):
|
||||||
|
"""register a function to be called when the stable diffusion model is created; the model is
|
||||||
|
passed as an argument"""
|
||||||
|
callbacks_model_loaded.append(callback)
|
||||||
|
|
||||||
|
|
||||||
|
def on_ui_tabs(callback):
|
||||||
|
"""register a function to be called when the UI is creating new tabs.
|
||||||
|
The function must either return a None, which means no new tabs to be added, or a list, where
|
||||||
|
each element is a tuple:
|
||||||
|
(gradio_component, title, elem_id)
|
||||||
|
|
||||||
|
gradio_component is a gradio component to be used for contents of the tab (usually gr.Blocks)
|
||||||
|
title is tab text displayed to user in the UI
|
||||||
|
elem_id is HTML id for the tab
|
||||||
|
"""
|
||||||
|
callbacks_ui_tabs.append(callback)
|
||||||
|
|
|
@ -1,86 +1,153 @@
|
||||||
import os
|
import os
|
||||||
import sys
|
import sys
|
||||||
import traceback
|
import traceback
|
||||||
|
from collections import namedtuple
|
||||||
|
|
||||||
import modules.ui as ui
|
import modules.ui as ui
|
||||||
import gradio as gr
|
import gradio as gr
|
||||||
|
|
||||||
from modules.processing import StableDiffusionProcessing
|
from modules.processing import StableDiffusionProcessing
|
||||||
from modules import shared
|
from modules import shared, paths, script_callbacks
|
||||||
|
|
||||||
|
AlwaysVisible = object()
|
||||||
|
|
||||||
|
|
||||||
class Script:
|
class Script:
|
||||||
filename = None
|
filename = None
|
||||||
args_from = None
|
args_from = None
|
||||||
args_to = None
|
args_to = None
|
||||||
|
alwayson = False
|
||||||
|
|
||||||
|
infotext_fields = None
|
||||||
|
"""if set in ui(), this is a list of pairs of gradio component + text; the text will be used when
|
||||||
|
parsing infotext to set the value for the component; see ui.py's txt2img_paste_fields for an example
|
||||||
|
"""
|
||||||
|
|
||||||
# The title of the script. This is what will be displayed in the dropdown menu.
|
|
||||||
def title(self):
|
def title(self):
|
||||||
|
"""this function should return the title of the script. This is what will be displayed in the dropdown menu."""
|
||||||
|
|
||||||
raise NotImplementedError()
|
raise NotImplementedError()
|
||||||
|
|
||||||
# How the script is displayed in the UI. See https://gradio.app/docs/#components
|
|
||||||
# for the different UI components you can use and how to create them.
|
|
||||||
# Most UI components can return a value, such as a boolean for a checkbox.
|
|
||||||
# The returned values are passed to the run method as parameters.
|
|
||||||
def ui(self, is_img2img):
|
def ui(self, is_img2img):
|
||||||
|
"""this function should create gradio UI elements. See https://gradio.app/docs/#components
|
||||||
|
The return value should be an array of all components that are used in processing.
|
||||||
|
Values of those returned componenbts will be passed to run() and process() functions.
|
||||||
|
"""
|
||||||
|
|
||||||
pass
|
pass
|
||||||
|
|
||||||
# Determines when the script should be shown in the dropdown menu via the
|
|
||||||
# returned value. As an example:
|
|
||||||
# is_img2img is True if the current tab is img2img, and False if it is txt2img.
|
|
||||||
# Thus, return is_img2img to only show the script on the img2img tab.
|
|
||||||
def show(self, is_img2img):
|
def show(self, is_img2img):
|
||||||
|
"""
|
||||||
|
is_img2img is True if this function is called for the img2img interface, and Fasle otherwise
|
||||||
|
|
||||||
|
This function should return:
|
||||||
|
- False if the script should not be shown in UI at all
|
||||||
|
- True if the script should be shown in UI if it's scelected in the scripts drowpdown
|
||||||
|
- script.AlwaysVisible if the script should be shown in UI at all times
|
||||||
|
"""
|
||||||
|
|
||||||
return True
|
return True
|
||||||
|
|
||||||
# This is where the additional processing is implemented. The parameters include
|
def run(self, p, *args):
|
||||||
# self, the model object "p" (a StableDiffusionProcessing class, see
|
"""
|
||||||
# processing.py), and the parameters returned by the ui method.
|
This function is called if the script has been selected in the script dropdown.
|
||||||
# Custom functions can be defined here, and additional libraries can be imported
|
It must do all processing and return the Processed object with results, same as
|
||||||
# to be used in processing. The return value should be a Processed object, which is
|
one returned by processing.process_images.
|
||||||
# what is returned by the process_images method.
|
|
||||||
def run(self, *args):
|
Usually the processing is done by calling the processing.process_images function.
|
||||||
|
|
||||||
|
args contains all values returned by components from ui()
|
||||||
|
"""
|
||||||
|
|
||||||
raise NotImplementedError()
|
raise NotImplementedError()
|
||||||
|
|
||||||
# The description method is currently unused.
|
def process(self, p, *args):
|
||||||
# To add a description that appears when hovering over the title, amend the "titles"
|
"""
|
||||||
# dict in script.js to include the script title (returned by title) as a key, and
|
This function is called before processing begins for AlwaysVisible scripts.
|
||||||
# your description as the value.
|
scripts. You can modify the processing object (p) here, inject hooks, etc.
|
||||||
|
"""
|
||||||
|
|
||||||
|
pass
|
||||||
|
|
||||||
def describe(self):
|
def describe(self):
|
||||||
|
"""unused"""
|
||||||
return ""
|
return ""
|
||||||
|
|
||||||
|
|
||||||
|
current_basedir = paths.script_path
|
||||||
|
|
||||||
|
|
||||||
|
def basedir():
|
||||||
|
"""returns the base directory for the current script. For scripts in the main scripts directory,
|
||||||
|
this is the main directory (where webui.py resides), and for scripts in extensions directory
|
||||||
|
(ie extensions/aesthetic/script/aesthetic.py), this is extension's directory (extensions/aesthetic)
|
||||||
|
"""
|
||||||
|
return current_basedir
|
||||||
|
|
||||||
|
|
||||||
scripts_data = []
|
scripts_data = []
|
||||||
|
ScriptFile = namedtuple("ScriptFile", ["basedir", "filename", "path"])
|
||||||
|
ScriptClassData = namedtuple("ScriptClassData", ["script_class", "path", "basedir"])
|
||||||
|
|
||||||
|
|
||||||
def load_scripts(basedir):
|
def list_scripts(scriptdirname, extension):
|
||||||
if not os.path.exists(basedir):
|
scripts_list = []
|
||||||
return
|
|
||||||
|
|
||||||
|
basedir = os.path.join(paths.script_path, scriptdirname)
|
||||||
|
if os.path.exists(basedir):
|
||||||
for filename in sorted(os.listdir(basedir)):
|
for filename in sorted(os.listdir(basedir)):
|
||||||
path = os.path.join(basedir, filename)
|
scripts_list.append(ScriptFile(paths.script_path, filename, os.path.join(basedir, filename)))
|
||||||
|
|
||||||
if os.path.splitext(path)[1].lower() != '.py':
|
extdir = os.path.join(paths.script_path, "extensions")
|
||||||
|
if os.path.exists(extdir):
|
||||||
|
for dirname in sorted(os.listdir(extdir)):
|
||||||
|
dirpath = os.path.join(extdir, dirname)
|
||||||
|
if not os.path.isdir(dirpath):
|
||||||
continue
|
continue
|
||||||
|
|
||||||
if not os.path.isfile(path):
|
for filename in sorted(os.listdir(os.path.join(dirpath, scriptdirname))):
|
||||||
continue
|
scripts_list.append(ScriptFile(dirpath, filename, os.path.join(dirpath, scriptdirname, filename)))
|
||||||
|
|
||||||
|
scripts_list = [x for x in scripts_list if os.path.splitext(x.path)[1].lower() == extension and os.path.isfile(x.path)]
|
||||||
|
|
||||||
|
return scripts_list
|
||||||
|
|
||||||
|
|
||||||
|
def load_scripts():
|
||||||
|
global current_basedir
|
||||||
|
scripts_data.clear()
|
||||||
|
script_callbacks.clear_callbacks()
|
||||||
|
|
||||||
|
scripts_list = list_scripts("scripts", ".py")
|
||||||
|
|
||||||
|
syspath = sys.path
|
||||||
|
|
||||||
|
for scriptfile in sorted(scripts_list):
|
||||||
try:
|
try:
|
||||||
with open(path, "r", encoding="utf8") as file:
|
if scriptfile.basedir != paths.script_path:
|
||||||
|
sys.path = [scriptfile.basedir] + sys.path
|
||||||
|
current_basedir = scriptfile.basedir
|
||||||
|
|
||||||
|
with open(scriptfile.path, "r", encoding="utf8") as file:
|
||||||
text = file.read()
|
text = file.read()
|
||||||
|
|
||||||
from types import ModuleType
|
from types import ModuleType
|
||||||
compiled = compile(text, path, 'exec')
|
compiled = compile(text, scriptfile.path, 'exec')
|
||||||
module = ModuleType(filename)
|
module = ModuleType(scriptfile.filename)
|
||||||
exec(compiled, module.__dict__)
|
exec(compiled, module.__dict__)
|
||||||
|
|
||||||
for key, script_class in module.__dict__.items():
|
for key, script_class in module.__dict__.items():
|
||||||
if type(script_class) == type and issubclass(script_class, Script):
|
if type(script_class) == type and issubclass(script_class, Script):
|
||||||
scripts_data.append((script_class, path))
|
scripts_data.append(ScriptClassData(script_class, scriptfile.path, scriptfile.basedir))
|
||||||
|
|
||||||
except Exception:
|
except Exception:
|
||||||
print(f"Error loading script: {filename}", file=sys.stderr)
|
print(f"Error loading script: {scriptfile.filename}", file=sys.stderr)
|
||||||
print(traceback.format_exc(), file=sys.stderr)
|
print(traceback.format_exc(), file=sys.stderr)
|
||||||
|
|
||||||
|
finally:
|
||||||
|
sys.path = syspath
|
||||||
|
current_basedir = paths.script_path
|
||||||
|
|
||||||
|
|
||||||
def wrap_call(func, filename, funcname, *args, default=None, **kwargs):
|
def wrap_call(func, filename, funcname, *args, default=None, **kwargs):
|
||||||
try:
|
try:
|
||||||
|
@ -96,56 +163,80 @@ def wrap_call(func, filename, funcname, *args, default=None, **kwargs):
|
||||||
class ScriptRunner:
|
class ScriptRunner:
|
||||||
def __init__(self):
|
def __init__(self):
|
||||||
self.scripts = []
|
self.scripts = []
|
||||||
|
self.selectable_scripts = []
|
||||||
|
self.alwayson_scripts = []
|
||||||
self.titles = []
|
self.titles = []
|
||||||
|
self.infotext_fields = []
|
||||||
|
|
||||||
def setup_ui(self, is_img2img):
|
def setup_ui(self, is_img2img):
|
||||||
for script_class, path in scripts_data:
|
for script_class, path, basedir in scripts_data:
|
||||||
script = script_class()
|
script = script_class()
|
||||||
script.filename = path
|
script.filename = path
|
||||||
|
|
||||||
if not script.show(is_img2img):
|
visibility = script.show(is_img2img)
|
||||||
continue
|
|
||||||
|
|
||||||
|
if visibility == AlwaysVisible:
|
||||||
self.scripts.append(script)
|
self.scripts.append(script)
|
||||||
|
self.alwayson_scripts.append(script)
|
||||||
|
script.alwayson = True
|
||||||
|
|
||||||
self.titles = [wrap_call(script.title, script.filename, "title") or f"{script.filename} [error]" for script in self.scripts]
|
elif visibility:
|
||||||
|
self.scripts.append(script)
|
||||||
|
self.selectable_scripts.append(script)
|
||||||
|
|
||||||
dropdown = gr.Dropdown(label="Script", choices=["None"] + self.titles, value="None", type="index")
|
self.titles = [wrap_call(script.title, script.filename, "title") or f"{script.filename} [error]" for script in self.selectable_scripts]
|
||||||
dropdown.save_to_config = True
|
|
||||||
inputs = [dropdown]
|
|
||||||
|
|
||||||
for script in self.scripts:
|
inputs = [None]
|
||||||
|
inputs_alwayson = [True]
|
||||||
|
|
||||||
|
def create_script_ui(script, inputs, inputs_alwayson):
|
||||||
script.args_from = len(inputs)
|
script.args_from = len(inputs)
|
||||||
script.args_to = len(inputs)
|
script.args_to = len(inputs)
|
||||||
|
|
||||||
controls = wrap_call(script.ui, script.filename, "ui", is_img2img)
|
controls = wrap_call(script.ui, script.filename, "ui", is_img2img)
|
||||||
|
|
||||||
if controls is None:
|
if controls is None:
|
||||||
continue
|
return
|
||||||
|
|
||||||
for control in controls:
|
for control in controls:
|
||||||
control.custom_script_source = os.path.basename(script.filename)
|
control.custom_script_source = os.path.basename(script.filename)
|
||||||
|
if not script.alwayson:
|
||||||
control.visible = False
|
control.visible = False
|
||||||
|
|
||||||
|
if script.infotext_fields is not None:
|
||||||
|
self.infotext_fields += script.infotext_fields
|
||||||
|
|
||||||
inputs += controls
|
inputs += controls
|
||||||
|
inputs_alwayson += [script.alwayson for _ in controls]
|
||||||
script.args_to = len(inputs)
|
script.args_to = len(inputs)
|
||||||
|
|
||||||
|
for script in self.alwayson_scripts:
|
||||||
|
with gr.Group():
|
||||||
|
create_script_ui(script, inputs, inputs_alwayson)
|
||||||
|
|
||||||
|
dropdown = gr.Dropdown(label="Script", choices=["None"] + self.titles, value="None", type="index")
|
||||||
|
dropdown.save_to_config = True
|
||||||
|
inputs[0] = dropdown
|
||||||
|
|
||||||
|
for script in self.selectable_scripts:
|
||||||
|
create_script_ui(script, inputs, inputs_alwayson)
|
||||||
|
|
||||||
def select_script(script_index):
|
def select_script(script_index):
|
||||||
if 0 < script_index <= len(self.scripts):
|
if 0 < script_index <= len(self.selectable_scripts):
|
||||||
script = self.scripts[script_index-1]
|
script = self.selectable_scripts[script_index-1]
|
||||||
args_from = script.args_from
|
args_from = script.args_from
|
||||||
args_to = script.args_to
|
args_to = script.args_to
|
||||||
else:
|
else:
|
||||||
args_from = 0
|
args_from = 0
|
||||||
args_to = 0
|
args_to = 0
|
||||||
|
|
||||||
return [ui.gr_show(True if i == 0 else args_from <= i < args_to) for i in range(len(inputs))]
|
return [ui.gr_show(True if i == 0 else args_from <= i < args_to or is_alwayson) for i, is_alwayson in enumerate(inputs_alwayson)]
|
||||||
|
|
||||||
def init_field(title):
|
def init_field(title):
|
||||||
if title == 'None':
|
if title == 'None':
|
||||||
return
|
return
|
||||||
script_index = self.titles.index(title)
|
script_index = self.titles.index(title)
|
||||||
script = self.scripts[script_index]
|
script = self.selectable_scripts[script_index]
|
||||||
for i in range(script.args_from, script.args_to):
|
for i in range(script.args_from, script.args_to):
|
||||||
inputs[i].visible = True
|
inputs[i].visible = True
|
||||||
|
|
||||||
|
@ -164,7 +255,7 @@ class ScriptRunner:
|
||||||
if script_index == 0:
|
if script_index == 0:
|
||||||
return None
|
return None
|
||||||
|
|
||||||
script = self.scripts[script_index-1]
|
script = self.selectable_scripts[script_index-1]
|
||||||
|
|
||||||
if script is None:
|
if script is None:
|
||||||
return None
|
return None
|
||||||
|
@ -176,6 +267,15 @@ class ScriptRunner:
|
||||||
|
|
||||||
return processed
|
return processed
|
||||||
|
|
||||||
|
def run_alwayson_scripts(self, p):
|
||||||
|
for script in self.alwayson_scripts:
|
||||||
|
try:
|
||||||
|
script_args = p.script_args[script.args_from:script.args_to]
|
||||||
|
script.process(p, *script_args)
|
||||||
|
except Exception:
|
||||||
|
print(f"Error running alwayson script: {script.filename}", file=sys.stderr)
|
||||||
|
print(traceback.format_exc(), file=sys.stderr)
|
||||||
|
|
||||||
def reload_sources(self):
|
def reload_sources(self):
|
||||||
for si, script in list(enumerate(self.scripts)):
|
for si, script in list(enumerate(self.scripts)):
|
||||||
with open(script.filename, "r", encoding="utf8") as file:
|
with open(script.filename, "r", encoding="utf8") as file:
|
||||||
|
@ -197,19 +297,21 @@ class ScriptRunner:
|
||||||
self.scripts[si].args_from = args_from
|
self.scripts[si].args_from = args_from
|
||||||
self.scripts[si].args_to = args_to
|
self.scripts[si].args_to = args_to
|
||||||
|
|
||||||
|
|
||||||
scripts_txt2img = ScriptRunner()
|
scripts_txt2img = ScriptRunner()
|
||||||
scripts_img2img = ScriptRunner()
|
scripts_img2img = ScriptRunner()
|
||||||
|
|
||||||
|
|
||||||
def reload_script_body_only():
|
def reload_script_body_only():
|
||||||
scripts_txt2img.reload_sources()
|
scripts_txt2img.reload_sources()
|
||||||
scripts_img2img.reload_sources()
|
scripts_img2img.reload_sources()
|
||||||
|
|
||||||
|
|
||||||
def reload_scripts(basedir):
|
def reload_scripts():
|
||||||
global scripts_txt2img, scripts_img2img
|
global scripts_txt2img, scripts_img2img
|
||||||
|
|
||||||
scripts_data.clear()
|
load_scripts()
|
||||||
load_scripts(basedir)
|
|
||||||
|
|
||||||
scripts_txt2img = ScriptRunner()
|
scripts_txt2img = ScriptRunner()
|
||||||
scripts_img2img = ScriptRunner()
|
scripts_img2img = ScriptRunner()
|
||||||
|
|
||||||
|
|
|
@ -332,7 +332,6 @@ class FrozenCLIPEmbedderWithCustomWords(torch.nn.Module):
|
||||||
multipliers.append([1.0] * 75)
|
multipliers.append([1.0] * 75)
|
||||||
|
|
||||||
z1 = self.process_tokens(tokens, multipliers)
|
z1 = self.process_tokens(tokens, multipliers)
|
||||||
z1 = shared.aesthetic_clip(z1, remade_batch_tokens)
|
|
||||||
z = z1 if z is None else torch.cat((z, z1), axis=-2)
|
z = z1 if z is None else torch.cat((z, z1), axis=-2)
|
||||||
|
|
||||||
remade_batch_tokens = rem_tokens
|
remade_batch_tokens = rem_tokens
|
||||||
|
|
|
@ -7,7 +7,7 @@ from omegaconf import OmegaConf
|
||||||
|
|
||||||
from ldm.util import instantiate_from_config
|
from ldm.util import instantiate_from_config
|
||||||
|
|
||||||
from modules import shared, modelloader, devices
|
from modules import shared, modelloader, devices, script_callbacks
|
||||||
from modules.paths import models_path
|
from modules.paths import models_path
|
||||||
from modules.sd_hijack_inpainting import do_inpainting_hijack, should_hijack_inpainting
|
from modules.sd_hijack_inpainting import do_inpainting_hijack, should_hijack_inpainting
|
||||||
|
|
||||||
|
@ -238,6 +238,9 @@ def load_model(checkpoint_info=None):
|
||||||
sd_hijack.model_hijack.hijack(sd_model)
|
sd_hijack.model_hijack.hijack(sd_model)
|
||||||
|
|
||||||
sd_model.eval()
|
sd_model.eval()
|
||||||
|
shared.sd_model = sd_model
|
||||||
|
|
||||||
|
script_callbacks.model_loaded_callback(sd_model)
|
||||||
|
|
||||||
print(f"Model loaded.")
|
print(f"Model loaded.")
|
||||||
return sd_model
|
return sd_model
|
||||||
|
@ -252,7 +255,7 @@ def reload_model_weights(sd_model, info=None):
|
||||||
|
|
||||||
if sd_model.sd_checkpoint_info.config != checkpoint_info.config or should_hijack_inpainting(checkpoint_info) != should_hijack_inpainting(sd_model.sd_checkpoint_info):
|
if sd_model.sd_checkpoint_info.config != checkpoint_info.config or should_hijack_inpainting(checkpoint_info) != should_hijack_inpainting(sd_model.sd_checkpoint_info):
|
||||||
checkpoints_loaded.clear()
|
checkpoints_loaded.clear()
|
||||||
shared.sd_model = load_model(checkpoint_info)
|
load_model(checkpoint_info)
|
||||||
return shared.sd_model
|
return shared.sd_model
|
||||||
|
|
||||||
if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
|
if shared.cmd_opts.lowvram or shared.cmd_opts.medvram:
|
||||||
|
|
|
@ -31,7 +31,6 @@ parser.add_argument("--no-half-vae", action='store_true', help="do not switch th
|
||||||
parser.add_argument("--no-progressbar-hiding", action='store_true', help="do not hide progressbar in gradio UI (we hide it because it slows down ML if you have hardware acceleration in browser)")
|
parser.add_argument("--no-progressbar-hiding", action='store_true', help="do not hide progressbar in gradio UI (we hide it because it slows down ML if you have hardware acceleration in browser)")
|
||||||
parser.add_argument("--max-batch-count", type=int, default=16, help="maximum batch count value for the UI")
|
parser.add_argument("--max-batch-count", type=int, default=16, help="maximum batch count value for the UI")
|
||||||
parser.add_argument("--embeddings-dir", type=str, default=os.path.join(script_path, 'embeddings'), help="embeddings directory for textual inversion (default: embeddings)")
|
parser.add_argument("--embeddings-dir", type=str, default=os.path.join(script_path, 'embeddings'), help="embeddings directory for textual inversion (default: embeddings)")
|
||||||
parser.add_argument("--aesthetic_embeddings-dir", type=str, default=os.path.join(models_path, 'aesthetic_embeddings'), help="aesthetic_embeddings directory(default: aesthetic_embeddings)")
|
|
||||||
parser.add_argument("--hypernetwork-dir", type=str, default=os.path.join(models_path, 'hypernetworks'), help="hypernetwork directory")
|
parser.add_argument("--hypernetwork-dir", type=str, default=os.path.join(models_path, 'hypernetworks'), help="hypernetwork directory")
|
||||||
parser.add_argument("--localizations-dir", type=str, default=os.path.join(script_path, 'localizations'), help="localizations directory")
|
parser.add_argument("--localizations-dir", type=str, default=os.path.join(script_path, 'localizations'), help="localizations directory")
|
||||||
parser.add_argument("--allow-code", action='store_true', help="allow custom script execution from webui")
|
parser.add_argument("--allow-code", action='store_true', help="allow custom script execution from webui")
|
||||||
|
@ -109,21 +108,6 @@ os.makedirs(cmd_opts.hypernetwork_dir, exist_ok=True)
|
||||||
hypernetworks = hypernetwork.list_hypernetworks(cmd_opts.hypernetwork_dir)
|
hypernetworks = hypernetwork.list_hypernetworks(cmd_opts.hypernetwork_dir)
|
||||||
loaded_hypernetwork = None
|
loaded_hypernetwork = None
|
||||||
|
|
||||||
|
|
||||||
os.makedirs(cmd_opts.aesthetic_embeddings_dir, exist_ok=True)
|
|
||||||
aesthetic_embeddings = {}
|
|
||||||
|
|
||||||
|
|
||||||
def update_aesthetic_embeddings():
|
|
||||||
global aesthetic_embeddings
|
|
||||||
aesthetic_embeddings = {f.replace(".pt", ""): os.path.join(cmd_opts.aesthetic_embeddings_dir, f) for f in
|
|
||||||
os.listdir(cmd_opts.aesthetic_embeddings_dir) if f.endswith(".pt")}
|
|
||||||
aesthetic_embeddings = OrderedDict(**{"None": None}, **aesthetic_embeddings)
|
|
||||||
|
|
||||||
|
|
||||||
update_aesthetic_embeddings()
|
|
||||||
|
|
||||||
|
|
||||||
def reload_hypernetworks():
|
def reload_hypernetworks():
|
||||||
global hypernetworks
|
global hypernetworks
|
||||||
|
|
||||||
|
@ -415,9 +399,6 @@ sd_model = None
|
||||||
|
|
||||||
clip_model = None
|
clip_model = None
|
||||||
|
|
||||||
from modules.aesthetic_clip import AestheticCLIP
|
|
||||||
aesthetic_clip = AestheticCLIP()
|
|
||||||
|
|
||||||
progress_print_out = sys.stdout
|
progress_print_out = sys.stdout
|
||||||
|
|
||||||
|
|
||||||
|
|
|
@ -7,7 +7,7 @@ import modules.processing as processing
|
||||||
from modules.ui import plaintext_to_html
|
from modules.ui import plaintext_to_html
|
||||||
|
|
||||||
|
|
||||||
def txt2img(prompt: str, negative_prompt: str, prompt_style: str, prompt_style2: str, steps: int, sampler_index: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, seed_enable_extras: bool, height: int, width: int, enable_hr: bool, denoising_strength: float, firstphase_width: int, firstphase_height: int, aesthetic_lr=0, aesthetic_weight=0, aesthetic_steps=0, aesthetic_imgs=None, aesthetic_slerp=False, aesthetic_imgs_text="", aesthetic_slerp_angle=0.15, aesthetic_text_negative=False, *args):
|
def txt2img(prompt: str, negative_prompt: str, prompt_style: str, prompt_style2: str, steps: int, sampler_index: int, restore_faces: bool, tiling: bool, n_iter: int, batch_size: int, cfg_scale: float, seed: int, subseed: int, subseed_strength: float, seed_resize_from_h: int, seed_resize_from_w: int, seed_enable_extras: bool, height: int, width: int, enable_hr: bool, denoising_strength: float, firstphase_width: int, firstphase_height: int, *args):
|
||||||
p = StableDiffusionProcessingTxt2Img(
|
p = StableDiffusionProcessingTxt2Img(
|
||||||
sd_model=shared.sd_model,
|
sd_model=shared.sd_model,
|
||||||
outpath_samples=opts.outdir_samples or opts.outdir_txt2img_samples,
|
outpath_samples=opts.outdir_samples or opts.outdir_txt2img_samples,
|
||||||
|
@ -36,7 +36,8 @@ def txt2img(prompt: str, negative_prompt: str, prompt_style: str, prompt_style2:
|
||||||
firstphase_height=firstphase_height if enable_hr else None,
|
firstphase_height=firstphase_height if enable_hr else None,
|
||||||
)
|
)
|
||||||
|
|
||||||
shared.aesthetic_clip.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)
|
p.scripts = modules.scripts.scripts_txt2img
|
||||||
|
p.script_args = args
|
||||||
|
|
||||||
if cmd_opts.enable_console_prompts:
|
if cmd_opts.enable_console_prompts:
|
||||||
print(f"\ntxt2img: {prompt}", file=shared.progress_print_out)
|
print(f"\ntxt2img: {prompt}", file=shared.progress_print_out)
|
||||||
|
|
|
@ -23,10 +23,10 @@ import gradio as gr
|
||||||
import gradio.utils
|
import gradio.utils
|
||||||
import gradio.routes
|
import gradio.routes
|
||||||
|
|
||||||
from modules import sd_hijack, sd_models, localization
|
from modules import sd_hijack, sd_models, localization, script_callbacks
|
||||||
from modules.paths import script_path
|
from modules.paths import script_path
|
||||||
|
|
||||||
from modules.shared import opts, cmd_opts, restricted_opts, aesthetic_embeddings
|
from modules.shared import opts, cmd_opts, restricted_opts
|
||||||
|
|
||||||
if cmd_opts.deepdanbooru:
|
if cmd_opts.deepdanbooru:
|
||||||
from modules.deepbooru import get_deepbooru_tags
|
from modules.deepbooru import get_deepbooru_tags
|
||||||
|
@ -44,7 +44,6 @@ from modules.images import save_image
|
||||||
import modules.textual_inversion.ui
|
import modules.textual_inversion.ui
|
||||||
import modules.hypernetworks.ui
|
import modules.hypernetworks.ui
|
||||||
|
|
||||||
import modules.aesthetic_clip as aesthetic_clip
|
|
||||||
import modules.images_history as img_his
|
import modules.images_history as img_his
|
||||||
|
|
||||||
|
|
||||||
|
@ -662,8 +661,6 @@ def create_ui(wrap_gradio_gpu_call):
|
||||||
|
|
||||||
seed, reuse_seed, subseed, reuse_subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w, seed_checkbox = create_seed_inputs()
|
seed, reuse_seed, subseed, reuse_subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w, seed_checkbox = create_seed_inputs()
|
||||||
|
|
||||||
aesthetic_weight, aesthetic_steps, aesthetic_lr, aesthetic_slerp, aesthetic_imgs, aesthetic_imgs_text, aesthetic_slerp_angle, aesthetic_text_negative = aesthetic_clip.create_ui()
|
|
||||||
|
|
||||||
with gr.Group():
|
with gr.Group():
|
||||||
custom_inputs = modules.scripts.scripts_txt2img.setup_ui(is_img2img=False)
|
custom_inputs = modules.scripts.scripts_txt2img.setup_ui(is_img2img=False)
|
||||||
|
|
||||||
|
@ -718,14 +715,6 @@ def create_ui(wrap_gradio_gpu_call):
|
||||||
denoising_strength,
|
denoising_strength,
|
||||||
firstphase_width,
|
firstphase_width,
|
||||||
firstphase_height,
|
firstphase_height,
|
||||||
aesthetic_lr,
|
|
||||||
aesthetic_weight,
|
|
||||||
aesthetic_steps,
|
|
||||||
aesthetic_imgs,
|
|
||||||
aesthetic_slerp,
|
|
||||||
aesthetic_imgs_text,
|
|
||||||
aesthetic_slerp_angle,
|
|
||||||
aesthetic_text_negative
|
|
||||||
] + custom_inputs,
|
] + custom_inputs,
|
||||||
|
|
||||||
outputs=[
|
outputs=[
|
||||||
|
@ -804,14 +793,7 @@ def create_ui(wrap_gradio_gpu_call):
|
||||||
(hr_options, lambda d: gr.Row.update(visible="Denoising strength" in d)),
|
(hr_options, lambda d: gr.Row.update(visible="Denoising strength" in d)),
|
||||||
(firstphase_width, "First pass size-1"),
|
(firstphase_width, "First pass size-1"),
|
||||||
(firstphase_height, "First pass size-2"),
|
(firstphase_height, "First pass size-2"),
|
||||||
(aesthetic_lr, "Aesthetic LR"),
|
*modules.scripts.scripts_txt2img.infotext_fields
|
||||||
(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"),
|
|
||||||
]
|
]
|
||||||
|
|
||||||
txt2img_preview_params = [
|
txt2img_preview_params = [
|
||||||
|
@ -896,8 +878,6 @@ def create_ui(wrap_gradio_gpu_call):
|
||||||
|
|
||||||
seed, reuse_seed, subseed, reuse_subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w, seed_checkbox = create_seed_inputs()
|
seed, reuse_seed, subseed, reuse_subseed, subseed_strength, seed_resize_from_h, seed_resize_from_w, seed_checkbox = create_seed_inputs()
|
||||||
|
|
||||||
aesthetic_weight_im, aesthetic_steps_im, aesthetic_lr_im, aesthetic_slerp_im, aesthetic_imgs_im, aesthetic_imgs_text_im, aesthetic_slerp_angle_im, aesthetic_text_negative_im = aesthetic_clip.create_ui()
|
|
||||||
|
|
||||||
with gr.Group():
|
with gr.Group():
|
||||||
custom_inputs = modules.scripts.scripts_img2img.setup_ui(is_img2img=True)
|
custom_inputs = modules.scripts.scripts_img2img.setup_ui(is_img2img=True)
|
||||||
|
|
||||||
|
@ -988,14 +968,6 @@ def create_ui(wrap_gradio_gpu_call):
|
||||||
inpainting_mask_invert,
|
inpainting_mask_invert,
|
||||||
img2img_batch_input_dir,
|
img2img_batch_input_dir,
|
||||||
img2img_batch_output_dir,
|
img2img_batch_output_dir,
|
||||||
aesthetic_lr_im,
|
|
||||||
aesthetic_weight_im,
|
|
||||||
aesthetic_steps_im,
|
|
||||||
aesthetic_imgs_im,
|
|
||||||
aesthetic_slerp_im,
|
|
||||||
aesthetic_imgs_text_im,
|
|
||||||
aesthetic_slerp_angle_im,
|
|
||||||
aesthetic_text_negative_im,
|
|
||||||
] + custom_inputs,
|
] + custom_inputs,
|
||||||
outputs=[
|
outputs=[
|
||||||
img2img_gallery,
|
img2img_gallery,
|
||||||
|
@ -1087,14 +1059,7 @@ def create_ui(wrap_gradio_gpu_call):
|
||||||
(seed_resize_from_w, "Seed resize from-1"),
|
(seed_resize_from_w, "Seed resize from-1"),
|
||||||
(seed_resize_from_h, "Seed resize from-2"),
|
(seed_resize_from_h, "Seed resize from-2"),
|
||||||
(denoising_strength, "Denoising strength"),
|
(denoising_strength, "Denoising strength"),
|
||||||
(aesthetic_lr_im, "Aesthetic LR"),
|
*modules.scripts.scripts_img2img.infotext_fields
|
||||||
(aesthetic_weight_im, "Aesthetic weight"),
|
|
||||||
(aesthetic_steps_im, "Aesthetic steps"),
|
|
||||||
(aesthetic_imgs_im, "Aesthetic embedding"),
|
|
||||||
(aesthetic_slerp_im, "Aesthetic slerp"),
|
|
||||||
(aesthetic_imgs_text_im, "Aesthetic text"),
|
|
||||||
(aesthetic_text_negative_im, "Aesthetic text negative"),
|
|
||||||
(aesthetic_slerp_angle_im, "Aesthetic slerp angle"),
|
|
||||||
]
|
]
|
||||||
token_button.click(fn=update_token_counter, inputs=[img2img_prompt, steps], outputs=[token_counter])
|
token_button.click(fn=update_token_counter, inputs=[img2img_prompt, steps], outputs=[token_counter])
|
||||||
|
|
||||||
|
@ -1217,9 +1182,9 @@ def create_ui(wrap_gradio_gpu_call):
|
||||||
)
|
)
|
||||||
#images history
|
#images history
|
||||||
images_history_switch_dict = {
|
images_history_switch_dict = {
|
||||||
"fn":modules.generation_parameters_copypaste.connect_paste,
|
"fn": modules.generation_parameters_copypaste.connect_paste,
|
||||||
"t2i":txt2img_paste_fields,
|
"t2i": txt2img_paste_fields,
|
||||||
"i2i":img2img_paste_fields
|
"i2i": img2img_paste_fields
|
||||||
}
|
}
|
||||||
|
|
||||||
images_history = img_his.create_history_tabs(gr, opts, cmd_opts, wrap_gradio_call(modules.extras.run_pnginfo), images_history_switch_dict)
|
images_history = img_his.create_history_tabs(gr, opts, cmd_opts, wrap_gradio_call(modules.extras.run_pnginfo), images_history_switch_dict)
|
||||||
|
@ -1264,18 +1229,6 @@ def create_ui(wrap_gradio_gpu_call):
|
||||||
with gr.Column():
|
with gr.Column():
|
||||||
create_embedding = gr.Button(value="Create embedding", variant='primary')
|
create_embedding = gr.Button(value="Create embedding", variant='primary')
|
||||||
|
|
||||||
with gr.Tab(label="Create aesthetic images embedding"):
|
|
||||||
|
|
||||||
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.Tab(label="Create hypernetwork"):
|
with gr.Tab(label="Create hypernetwork"):
|
||||||
new_hypernetwork_name = gr.Textbox(label="Name")
|
new_hypernetwork_name = gr.Textbox(label="Name")
|
||||||
new_hypernetwork_sizes = gr.CheckboxGroup(label="Modules", value=["768", "320", "640", "1280"], choices=["768", "320", "640", "1280"])
|
new_hypernetwork_sizes = gr.CheckboxGroup(label="Modules", value=["768", "320", "640", "1280"], choices=["768", "320", "640", "1280"])
|
||||||
|
@ -1375,21 +1328,6 @@ def create_ui(wrap_gradio_gpu_call):
|
||||||
]
|
]
|
||||||
)
|
)
|
||||||
|
|
||||||
create_embedding_ae.click(
|
|
||||||
fn=aesthetic_clip.generate_imgs_embd,
|
|
||||||
inputs=[
|
|
||||||
new_embedding_name_ae,
|
|
||||||
process_src_ae,
|
|
||||||
batch_ae
|
|
||||||
],
|
|
||||||
outputs=[
|
|
||||||
aesthetic_imgs,
|
|
||||||
aesthetic_imgs_im,
|
|
||||||
ti_output,
|
|
||||||
ti_outcome,
|
|
||||||
]
|
|
||||||
)
|
|
||||||
|
|
||||||
create_hypernetwork.click(
|
create_hypernetwork.click(
|
||||||
fn=modules.hypernetworks.ui.create_hypernetwork,
|
fn=modules.hypernetworks.ui.create_hypernetwork,
|
||||||
inputs=[
|
inputs=[
|
||||||
|
@ -1580,10 +1518,10 @@ Requested path was: {f}
|
||||||
if not opts.same_type(value, opts.data_labels[key].default):
|
if not opts.same_type(value, opts.data_labels[key].default):
|
||||||
return gr.update(visible=True), opts.dumpjson()
|
return gr.update(visible=True), opts.dumpjson()
|
||||||
|
|
||||||
|
oldval = opts.data.get(key, None)
|
||||||
if cmd_opts.hide_ui_dir_config and key in restricted_opts:
|
if cmd_opts.hide_ui_dir_config and key in restricted_opts:
|
||||||
return gr.update(value=oldval), opts.dumpjson()
|
return gr.update(value=oldval), opts.dumpjson()
|
||||||
|
|
||||||
oldval = opts.data.get(key, None)
|
|
||||||
opts.data[key] = value
|
opts.data[key] = value
|
||||||
|
|
||||||
if oldval != value:
|
if oldval != value:
|
||||||
|
@ -1692,9 +1630,12 @@ Requested path was: {f}
|
||||||
(images_history, "Image Browser", "images_history"),
|
(images_history, "Image Browser", "images_history"),
|
||||||
(modelmerger_interface, "Checkpoint Merger", "modelmerger"),
|
(modelmerger_interface, "Checkpoint Merger", "modelmerger"),
|
||||||
(train_interface, "Train", "ti"),
|
(train_interface, "Train", "ti"),
|
||||||
(settings_interface, "Settings", "settings"),
|
|
||||||
]
|
]
|
||||||
|
|
||||||
|
interfaces += script_callbacks.ui_tabs_callback()
|
||||||
|
|
||||||
|
interfaces += [(settings_interface, "Settings", "settings")]
|
||||||
|
|
||||||
with open(os.path.join(script_path, "style.css"), "r", encoding="utf8") as file:
|
with open(os.path.join(script_path, "style.css"), "r", encoding="utf8") as file:
|
||||||
css = file.read()
|
css = file.read()
|
||||||
|
|
||||||
|
|
7
webui.py
7
webui.py
|
@ -71,6 +71,7 @@ def wrap_gradio_gpu_call(func, extra_outputs=None):
|
||||||
|
|
||||||
return modules.ui.wrap_gradio_call(f, extra_outputs=extra_outputs)
|
return modules.ui.wrap_gradio_call(f, extra_outputs=extra_outputs)
|
||||||
|
|
||||||
|
|
||||||
def initialize():
|
def initialize():
|
||||||
modelloader.cleanup_models()
|
modelloader.cleanup_models()
|
||||||
modules.sd_models.setup_model()
|
modules.sd_models.setup_model()
|
||||||
|
@ -79,9 +80,9 @@ def initialize():
|
||||||
shared.face_restorers.append(modules.face_restoration.FaceRestoration())
|
shared.face_restorers.append(modules.face_restoration.FaceRestoration())
|
||||||
modelloader.load_upscalers()
|
modelloader.load_upscalers()
|
||||||
|
|
||||||
modules.scripts.load_scripts(os.path.join(script_path, "scripts"))
|
modules.scripts.load_scripts()
|
||||||
|
|
||||||
shared.sd_model = modules.sd_models.load_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)))
|
shared.opts.onchange("sd_model_checkpoint", wrap_queued_call(lambda: modules.sd_models.reload_model_weights(shared.sd_model)))
|
||||||
shared.opts.onchange("sd_hypernetwork", wrap_queued_call(lambda: modules.hypernetworks.hypernetwork.load_hypernetwork(shared.opts.sd_hypernetwork)))
|
shared.opts.onchange("sd_hypernetwork", wrap_queued_call(lambda: modules.hypernetworks.hypernetwork.load_hypernetwork(shared.opts.sd_hypernetwork)))
|
||||||
shared.opts.onchange("sd_hypernetwork_strength", modules.hypernetworks.hypernetwork.apply_strength)
|
shared.opts.onchange("sd_hypernetwork_strength", modules.hypernetworks.hypernetwork.apply_strength)
|
||||||
|
@ -145,7 +146,7 @@ def webui():
|
||||||
sd_samplers.set_samplers()
|
sd_samplers.set_samplers()
|
||||||
|
|
||||||
print('Reloading Custom Scripts')
|
print('Reloading Custom Scripts')
|
||||||
modules.scripts.reload_scripts(os.path.join(script_path, "scripts"))
|
modules.scripts.reload_scripts()
|
||||||
print('Reloading modules: modules.ui')
|
print('Reloading modules: modules.ui')
|
||||||
importlib.reload(modules.ui)
|
importlib.reload(modules.ui)
|
||||||
print('Refreshing Model List')
|
print('Refreshing Model List')
|
||||||
|
|
Loading…
Reference in a new issue