├── embedding_editor_v1.png ├── embedding_editor_v2.png ├── style.css ├── javascript └── embedding_editor.js ├── README.md └── scripts └── embedding_editor.py /embedding_editor_v1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CodeExplode/stable-diffusion-webui-embedding-editor/HEAD/embedding_editor_v1.png -------------------------------------------------------------------------------- /embedding_editor_v2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/CodeExplode/stable-diffusion-webui-embedding-editor/HEAD/embedding_editor_v2.png -------------------------------------------------------------------------------- /style.css: -------------------------------------------------------------------------------- 1 | #embedding_editor_weight_sliders_container{ 2 | height: 745px; 3 | overflow-y: auto !important; 4 | resize: vertical; 5 | padding-top: 7px; 6 | } 7 | 8 | #embedding_editor_refresh_guidance{ 9 | max-width: 2em; 10 | } 11 | 12 | [id^="embedding_editor_weight_slider_"] input[type=range] { 13 | 14 | } -------------------------------------------------------------------------------- /javascript/embedding_editor.js: -------------------------------------------------------------------------------- 1 | function embedding_editor_update_guidance(col_weights) { 2 | let weightsColsVals = {}; 3 | for (let i=0; i<768; i++) 4 | weightsColsVals[i] = []; 5 | 6 | for (let color in col_weights) { 7 | let weights = Object.values(col_weights[color]); 8 | 9 | for (let i=0; i<768; i++) 10 | embedding_editor_sorted_insert(weightsColsVals[i], color, weights[i]); 11 | } 12 | 13 | for (let i=0; i<768; i++) { 14 | let guidanceElem = embedding_editor_get_guidance_bar(i), 15 | variants = weightsColsVals[i]; 16 | 17 | // e.g. linear-gradient(to right, transparent 0%, transparent 31%, red 31%, red 34%, transparent 34%, transparent 100%) 18 | let transparent = '#00000000', 19 | bg = 'linear-gradient(to right', 20 | previous = 0; 21 | 22 | for (let j=0, jLen=variants.length; j weight) 50 | break; 51 | 52 | set.splice(insertAt, 0, { color : color, weight: weight }); 53 | } 54 | 55 | function embedding_editor_get_guidance_bar(weightNum) { 56 | let guidanceClassType = "embedding_editor_guidance_bar", 57 | gradioSlider = document.querySelector("gradio-app").shadowRoot.getElementById("embedding_editor_weight_slider_" + weightNum), 58 | childElems = Array.from(gradioSlider.childNodes), 59 | lastChild = childElems[childElems.length-1]; 60 | 61 | if (!lastChild.classList.contains(guidanceClassType)) { 62 | // currently pointing to the range input. Move it slightly lower to line up with the new div, then insert the new div 63 | lastChild.style.verticalAlign = 'text-bottom'; // could do this with CSS, have the selector for it, and then won't change on first time pressing button 64 | 65 | let newElem = document.createElement("div"); 66 | newElem.style.height='6px'; 67 | newElem.classList.add(guidanceClassType); 68 | gradioSlider.appendChild(newElem); 69 | lastChild = newElem; 70 | } 71 | 72 | return lastChild; // could just cache these 73 | } 74 | 75 | onUiUpdate(function(){ 76 | 77 | }) 78 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ![preview image](embedding_editor_v2.png) 2 | 3 | A very early WIP of an embeddings editor for AUTOMATIC1111's webui. 4 | 5 | Should be placed in a sub-directory in extensions. e.g. \stable-diffusion-webui-master\extensions\embedding-editor\ You can now also install it automatically from the extensions tab in Automatic's UI. 6 | 7 | It will likely add a small amount of startup time due to fetching all the original embeddings to calculate the ranges for the weight sliders, so it's probably better to only enable when you're using it for now. 8 | 9 | My hope is that instead of say typing 'red shirt' and 'blue castle' and having those colours bleed into everything, you'd create a quick embedding RedShirt by starting with 'shirt' and shifting some weights to get an embedding for red shirts only, since a lot of embeddings have colour information which doesn't bleed into everything else in the image (e.g. apples, skin, trees, skies, water). 10 | 11 | I'm not sure yet how possible that would be, but you can literally just blend the puppy and skunk embeddings (with 50% of each) to get a very functional SkunkPuppy embedding, like https://www.youtube.com/watch?v=0_BBRNYInx8&t=888s , which is much more useful than trying to use both skunk and puppy in a prompt. So I'm hoping that with the ability to tag the positions of other embeddings on the weight sliders, you can get an idea of what can and can't be shifted to stay within the bounds of difference concepts, and we might be able to figure out if each weight maps to particular features which can be labelled. 12 | 13 | In my experience playing with it a bit so far, changing the vast majority of weights barely changes the output at all, which makes me hopeful that a lot of them are actually junk weights (or only apply for certain types of things which are unlocked by other weights), and just a few need to be paid attention to. 14 | 15 | It currently doesn't check for invalid inputs and will likely error easily (e.g. you can enter invalid vector numbers, or click save embedding while no embedding is selected). 16 | 17 | 'Refresh Embeddings' will repopulate the embeddings list, though won't reset the weight sliders on the currently selected embedding. Select a different embedding or change the vector number to refresh the weight sliders. 18 | 19 | Changing the vector number will lose any edited weights for the current vector. Saving the embedding first is required if you want to keep them. 20 | 21 | It's not easy to get any useful results with this yet, though some features are planned which might help a lot. 22 | 23 | These embeddings can be used as a starting point for textual inversion, or the results of textual inversion can be manually fine-tuned using these sliders. There are multiple valid embeddings for a given subject, so keeping certain weights in a range which causes the embedding to behave as the type of object which you want (e.g. a shirt, a person, a pose) may be useful, though would be better worked into the textual inversion process as a regulation method. 24 | 25 | --- 26 | 27 | I'd also recommend checking out Embedding Inspector, which looks like it does a lot of what I was intending to do: https://github.com/tkalayci71/embedding-inspector 28 | -------------------------------------------------------------------------------- /scripts/embedding_editor.py: -------------------------------------------------------------------------------- 1 | import os 2 | from webui import wrap_gradio_gpu_call 3 | from modules import scripts, script_callbacks 4 | from modules import shared, devices, sd_hijack, processing, sd_models, images, ui 5 | from modules.shared import opts, cmd_opts, restricted_opts 6 | from modules.ui import create_output_panel, setup_progressbar, create_refresh_button 7 | from modules.processing import StableDiffusionProcessing, Processed, StableDiffusionProcessingTxt2Img, \ 8 | StableDiffusionProcessingImg2Img, process_images 9 | from modules.ui import plaintext_to_html 10 | from modules.textual_inversion.textual_inversion import save_embedding 11 | import gradio as gr 12 | import gradio.routes 13 | import gradio.utils 14 | import torch 15 | 16 | # ISSUES 17 | # distribution shouldn't be fetched until the first embedding is opened, and can probably be converted into a numpy array 18 | # most functions need to verify that an embedding is selected 19 | # vector numbers aren't verified (might be better as a slider) 20 | # weight slider values are lost when changing vector number 21 | # remove unused imports 22 | # 23 | # TODO 24 | # add tagged positions on sliders from user-supplied words (and unique symbols & colours) 25 | # add a word->substrings printout for use with the above for words which map to multiple embeddings (e.g. "computer" = "compu" and "ter") 26 | # add the ability to create embeddings which are a mix of other embeddings (with ratios), e.g. 0.5 * skunk + 0.5 * puppy is a valid embedding 27 | # add the ability to shift all weights towards another embedding with a master slider 28 | # add a strength slider (multiply all weights) 29 | # print out the closest word(s) in the original embeddings list to the current embedding, with torch.abs(embedding1.vec - embedding2.vec).mean() or maybe sum 30 | # also maybe print a mouseover or have an expandable per weight slider for the closest embedding(s) for that weight value 31 | # maybe allowing per-weight notes, and possibly a way to save them per embedding vector 32 | # add option to vary individual weights one at a time and geneerate outputs, potentially also combinations of weights. Potentially use scoring system to determine size of change (maybe latents or clip interrogator) 33 | # add option to 'move' around current embedding position and generate outputs (a 768-dimensional vector spiral)? 34 | 35 | embedding_editor_weight_visual_scalar = 1 36 | 37 | def determine_embedding_distribution(): 38 | cond_model = shared.sd_model.cond_stage_model 39 | embedding_layer = cond_model.wrapped.transformer.text_model.embeddings 40 | 41 | # fix for medvram/lowvram - can't figure out how to detect the device of the model in torch, so will try to guess from the web ui options 42 | device = devices.device 43 | if cmd_opts.medvram or cmd_opts.lowvram: 44 | device = torch.device("cpu") 45 | # 46 | 47 | for i in range(49405): # guessing that's the range of CLIP tokens given that 49406 and 49407 are special tokens presumably appended to the end 48 | embedding = embedding_layer.token_embedding.wrapped(torch.LongTensor([i]).to(device)).squeeze(0) 49 | if i == 0: 50 | distribution_floor = embedding 51 | distribution_ceiling = embedding 52 | else: 53 | distribution_floor = torch.minimum(distribution_floor, embedding) 54 | distribution_ceiling = torch.maximum(distribution_ceiling, embedding) 55 | 56 | # a hack but don't know how else to get these values into gradio event functions, short of maybe caching them in an invisible gradio html element 57 | global embedding_editor_distribution_floor, embedding_editor_distribution_ceiling 58 | embedding_editor_distribution_floor = distribution_floor 59 | embedding_editor_distribution_ceiling = distribution_ceiling 60 | 61 | def build_slider(index, default, weight_sliders): 62 | floor = embedding_editor_distribution_floor[index].item() * embedding_editor_weight_visual_scalar 63 | ceil = embedding_editor_distribution_ceiling[index].item() * embedding_editor_weight_visual_scalar 64 | 65 | slider = gr.Slider(minimum=floor, maximum=ceil, step="any", label=f"w{index}", value=default, interactive=True, elem_id=f'embedding_editor_weight_slider_{index}') 66 | 67 | weight_sliders.append(slider) 68 | 69 | def on_ui_tabs(): 70 | determine_embedding_distribution() 71 | weight_sliders = [] 72 | 73 | with gr.Blocks(analytics_enabled=False) as embedding_editor_interface: 74 | with gr.Row().style(equal_height=False): 75 | with gr.Column(variant='panel', scale=1.5): 76 | with gr.Column(): 77 | with gr.Row(): 78 | embedding_name = gr.Dropdown(label='Embedding', elem_id="edit_embedding", choices=sorted(sd_hijack.model_hijack.embedding_db.word_embeddings.keys()), interactive=True) 79 | vector_num = gr.Number(label='Vector', value=0, step=1, interactive=True) 80 | refresh_embeddings_button = gr.Button(value="Refresh Embeddings", variant='secondary') 81 | save_embedding_button = gr.Button(value="Save Embedding", variant='primary') 82 | 83 | instructions = gr.HTML(f""" 84 |

Enter words and color hexes to mark weights on the sliders for guidance. Hint: Use the txt2img prompt token counter or webui-tokenizer to see which words are constructed using multiple sub-words, e.g. 'computer' doesn't exist in stable diffusion's CLIP dictionary and instead 'compu' and 'ter' are used (1 word but 2 embedding vectors). Currently buggy and needs a moment to process before pressing the button. If it doesn't work after a moment, try adding a random space to refresh it. 85 |

86 | """) 87 | with gr.Row(): 88 | guidance_embeddings = gr.Textbox(value="apple:#FF0000, banana:#FECE26, strawberry:#FF00FF", placeholder="symbol:color-hex, symbol:color-hex, ...", show_label=False, interactive=True) 89 | guidance_update_button = gr.Button(value='\U0001f504', elem_id='embedding_editor_refresh_guidance') 90 | guidance_hidden_cache = gr.HTML(value="", visible=False) 91 | 92 | with gr.Column(elem_id='embedding_editor_weight_sliders_container'): 93 | for i in range(0, 128): 94 | with gr.Row(): 95 | build_slider(i*6+0, 0, weight_sliders) 96 | build_slider(i*6+1, 0, weight_sliders) 97 | build_slider(i*6+2, 0, weight_sliders) 98 | build_slider(i*6+3, 0, weight_sliders) 99 | build_slider(i*6+4, 0, weight_sliders) 100 | build_slider(i*6+5, 0, weight_sliders) 101 | 102 | with gr.Column(scale=1): 103 | gallery = gr.Gallery(label='Output', show_label=False, elem_id="embedding_editor_gallery").style(grid=4) 104 | prompt = gr.Textbox(label="Prompt", elem_id=f"embedding_editor_prompt", show_label=False, lines=2, placeholder="e.g. A portrait photo of embedding_name" ) 105 | batch_count = gr.Slider(minimum=1, step=1, label='Batch count', value=1) 106 | steps = gr.Slider(minimum=1, maximum=150, step=1, label="Sampling Steps", value=20) 107 | cfg_scale = gr.Slider(minimum=1.0, maximum=30.0, step=0.5, label='CFG Scale', value=7.0) 108 | seed =(gr.Textbox if cmd_opts.use_textbox_seed else gr.Number)(label='Seed', value=-1) 109 | 110 | with gr.Row(): 111 | generate_preview = gr.Button(value="Generate Preview", variant='primary') 112 | 113 | generation_info = gr.HTML() 114 | html_info = gr.HTML() 115 | 116 | preview_args = dict( 117 | fn=wrap_gradio_gpu_call(generate_embedding_preview), 118 | #_js="submit", 119 | inputs=[ 120 | embedding_name, 121 | vector_num, 122 | prompt, 123 | steps, 124 | cfg_scale, 125 | seed, 126 | batch_count, 127 | ] + weight_sliders, 128 | outputs=[ 129 | gallery, 130 | generation_info, 131 | html_info 132 | ], 133 | show_progress=False, 134 | ) 135 | 136 | generate_preview.click(**preview_args) 137 | 138 | selection_args = dict( 139 | fn=select_embedding, 140 | inputs=[ 141 | embedding_name, 142 | vector_num, 143 | ], 144 | outputs = weight_sliders, 145 | ) 146 | 147 | embedding_name.change(**selection_args) 148 | vector_num.change(**selection_args) 149 | 150 | def refresh_embeddings(): 151 | sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings() # refresh_method 152 | 153 | refreshed_args = lambda: {"choices": sorted(sd_hijack.model_hijack.embedding_db.word_embeddings.keys())} # refreshed_args 154 | args = refreshed_args() if callable(refreshed_args) else refreshed_args 155 | 156 | for k, v in args.items(): 157 | setattr(embedding_name, k, v) 158 | 159 | return gr.update(**(args or {})) 160 | 161 | refresh_embeddings_button.click( 162 | fn=refresh_embeddings, 163 | inputs=[], 164 | outputs=[embedding_name] 165 | ) 166 | 167 | save_embedding_button.click( 168 | fn=save_embedding_weights, 169 | inputs=[ 170 | embedding_name, 171 | vector_num, 172 | ] + weight_sliders, 173 | outputs=[], 174 | ) 175 | 176 | guidance_embeddings.change( 177 | fn=update_guidance_embeddings, 178 | inputs=[guidance_embeddings], 179 | outputs=[guidance_hidden_cache] 180 | ) 181 | 182 | guidance_update_button.click( 183 | fn=None, 184 | _js="embedding_editor_update_guidance", 185 | inputs=[guidance_hidden_cache], 186 | outputs=[] 187 | ) 188 | 189 | guidance_hidden_cache.value = update_guidance_embeddings(guidance_embeddings.value) 190 | 191 | return [(embedding_editor_interface, "Embedding Editor", "embedding_editor_interface")] 192 | 193 | def select_embedding(embedding_name, vector_num): 194 | embedding = sd_hijack.model_hijack.embedding_db.word_embeddings[embedding_name] 195 | vec = embedding.vec[int(vector_num)] 196 | weights = [] 197 | 198 | for i in range(0, 768): 199 | weights.append( vec[i].item() * embedding_editor_weight_visual_scalar ) 200 | 201 | return weights 202 | 203 | def apply_slider_weights(embedding_name, vector_num, weights): 204 | embedding = sd_hijack.model_hijack.embedding_db.word_embeddings[embedding_name] 205 | vec = embedding.vec[int(vector_num)] 206 | old_weights = [] 207 | 208 | for i in range(0, 768): 209 | old_weights.append(vec[i].item()) 210 | vec[i] = weights[i] / embedding_editor_weight_visual_scalar 211 | 212 | return old_weights 213 | 214 | def generate_embedding_preview(embedding_name, vector_num, prompt: str, steps: int, cfg_scale: float, seed: int, batch_count: int, *weights): 215 | old_weights = apply_slider_weights(embedding_name, vector_num, weights) 216 | 217 | p = StableDiffusionProcessingTxt2Img( 218 | sd_model=shared.sd_model, 219 | outpath_samples=opts.outdir_samples or opts.outdir_txt2img_samples, 220 | outpath_grids=opts.outdir_grids or opts.outdir_txt2img_grids, 221 | prompt=prompt, 222 | seed=seed, 223 | steps=steps, 224 | cfg_scale=cfg_scale, 225 | n_iter=batch_count, 226 | ) 227 | 228 | if cmd_opts.enable_console_prompts: 229 | print(f"\ntxt2img: {prompt}", file=shared.progress_print_out) 230 | 231 | processed = process_images(p) 232 | 233 | p.close() 234 | 235 | shared.total_tqdm.clear() 236 | 237 | generation_info_js = processed.js() 238 | if opts.samples_log_stdout: 239 | print(generation_info_js) 240 | 241 | apply_slider_weights(embedding_name, vector_num, old_weights) # restore 242 | 243 | return processed.images, generation_info_js, plaintext_to_html(processed.info) 244 | 245 | def save_embedding_weights(embedding_name, vector_num, *weights): 246 | apply_slider_weights(embedding_name, vector_num, weights) 247 | embedding = sd_hijack.model_hijack.embedding_db.word_embeddings[embedding_name] 248 | checkpoint = sd_models.select_checkpoint() 249 | 250 | filename = os.path.join(shared.cmd_opts.embeddings_dir, f'{embedding_name}.pt') 251 | save_embedding(embedding, checkpoint, embedding_name, filename, remove_cached_checksum=True) 252 | 253 | def update_guidance_embeddings(text): 254 | try: 255 | cond_model = shared.sd_model.cond_stage_model 256 | embedding_layer = cond_model.wrapped.transformer.text_model.embeddings 257 | 258 | pairs = [x.strip() for x in text.split(',')] 259 | 260 | col_weights = {} 261 | 262 | for pair in pairs: 263 | word, col = pair.split(":") 264 | 265 | ids = cond_model.tokenizer(word, max_length=77, return_tensors="pt", add_special_tokens=False)["input_ids"] 266 | embedding = embedding_layer.token_embedding.wrapped(ids.to(devices.device)).squeeze(0)[0] 267 | weights = [] 268 | 269 | for i in range(0, 768): 270 | weight = embedding[i].item() 271 | floor = embedding_editor_distribution_floor[i].item() 272 | ceiling = embedding_editor_distribution_ceiling[i].item() 273 | 274 | weight = (weight - floor) / (ceiling - floor) # adjust to range for using as a guidance marker along the slider 275 | weights.append(weight) 276 | 277 | col_weights[col] = weights 278 | 279 | return col_weights 280 | except: 281 | return [] 282 | 283 | 284 | script_callbacks.on_ui_tabs(on_ui_tabs) 285 | --------------------------------------------------------------------------------