├── README.md └── embedding_to_png.py /README.md: -------------------------------------------------------------------------------- 1 | # embedding-to-png-script 2 | 3 | Script for https://github.com/AUTOMATIC1111/stable-diffusion-webui to convert .pt embeddings to safe PNGs. 4 | 5 | ## Parameters 6 | 7 | - *Source embedding to convert* - an existing embedding file to convert. 8 | - *Embedding token* - toke to use for the embedding and display on the image. 9 | - *Output directory* - Output folder. 10 | 11 | You can leave the prompt and **Embedding token** blank and the cleaned filename from *Source embedding to convert* will be used for both. 12 | 13 | ## Output examples 14 | 15 | ![spritual-monsters](https://user-images.githubusercontent.com/35278260/196052398-268a3a3e-0fad-46cd-b37d-9808480ceb18.png) 16 | ![ricar](https://user-images.githubusercontent.com/35278260/196052411-2c25ff22-407b-487b-a8c7-d04d712816f7.png) 17 | ![dr-strange](https://user-images.githubusercontent.com/35278260/196052418-b232ee90-11f9-4c32-8091-f52e0b3b009f.png) 18 | -------------------------------------------------------------------------------- /embedding_to_png.py: -------------------------------------------------------------------------------- 1 | import modules.scripts as scripts 2 | import gradio as gr 3 | import os 4 | import torch 5 | import io 6 | from modules import devices, sd_models, sd_hijack 7 | from PIL import PngImagePlugin 8 | from modules.processing import Processed, process_images 9 | from modules.shared import opts, cmd_opts, state 10 | from modules.textual_inversion.image_embedding import ( 11 | caption_image_overlay, 12 | insert_image_data_embed, 13 | extract_image_data_embed, 14 | embedding_to_b64, 15 | embedding_from_b64, 16 | ) 17 | 18 | 19 | class Script(scripts.Script): 20 | 21 | def title(self): 22 | return "Embedding to Shareable PNG" 23 | 24 | def ui(self, is_img2img): 25 | embedding = gr.File(label="Source embedding to convert") 26 | embedding_token = gr.Textbox(label="Embedding token") 27 | destination_folder = gr.Textbox(label="Output directory", value="outputs") 28 | return [embedding, embedding_token, destination_folder] 29 | 30 | def run(self, p, embedding, embedding_token, destination_folder): 31 | print(embedding, embedding_token, destination_folder) 32 | assert os.path.exists(destination_folder) 33 | sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings() 34 | 35 | embedding_token = embedding_token.replace('<', '').replace('>', '').strip() 36 | 37 | try: 38 | data = torch.load(embedding.name) 39 | except Exception as e: 40 | data = extract_image_data_embed(Image.open(embedding.name)) 41 | 42 | assert data is not None 43 | 44 | original_name = os.path.splitext(os.path.basename(os.path.split(embedding.orig_name)[-1]))[0] 45 | 46 | if p.prompt == '': 47 | p.prompt = original_name 48 | 49 | if embedding_token == '': 50 | embedding_token = original_name 51 | 52 | # textual inversion embeddings 53 | if 'string_to_param' in data: 54 | param_dict = data['string_to_param'] 55 | if hasattr(param_dict, '_parameters'): 56 | param_dict = getattr(param_dict, '_parameters') # fix for torch 1.12.1 loading saved file from torch 1.11 57 | assert len(param_dict) == 1, 'embedding file has multiple terms in it' 58 | emb = next(iter(param_dict.items()))[1] 59 | # diffuser concepts 60 | elif type(data) == dict and type(next(iter(data.values()))) == torch.Tensor: 61 | assert len(data.keys()) == 1, 'embedding file has multiple terms in it' 62 | 63 | emb = next(iter(data.values())) 64 | if len(emb.shape) == 1: 65 | emb = emb.unsqueeze(0) 66 | else: 67 | raise Exception(f"Couldn't identify embedding as either textual inversion embedding nor diffuser concept.") 68 | 69 | checkpoint = sd_models.select_checkpoint() 70 | 71 | emb_data = { 72 | "string_to_token": {"*": 265}, 73 | "string_to_param": {"*": emb.detach().to(devices.device, dtype=torch.float32)}, 74 | "name": embedding_token, 75 | "step": data.get('step', 0), 76 | "sd_checkpoint": data.get('hash', None), 77 | "sd_checkpoint_name": data.get('sd_checkpoint_name', None), 78 | } 79 | 80 | data = emb_data 81 | 82 | processed = process_images(p) 83 | image = processed.images[0] 84 | 85 | title = ' ' 86 | 87 | if 'name' in data: 88 | title = "<{}>".format(embedding_token) 89 | 90 | info = PngImagePlugin.PngInfo() 91 | data['name'] = embedding_token 92 | info.add_text("sd-ti-embedding", embedding_to_b64(data)) 93 | 94 | try: 95 | vectorSize = list(data['string_to_param'].values())[0].shape[0] 96 | except Exception as e: 97 | vectorSize = None 98 | 99 | footer_left = checkpoint.model_name 100 | footer_mid = '[{}]'.format(checkpoint.hash) 101 | footer_right = ' ' 102 | 103 | if vectorSize is not None: 104 | footer_right += '{}v'.format(vectorSize) 105 | 106 | if data.get('step', 0) > 0: 107 | footer_right += ' {}s'.format(data.get('step', 0)) 108 | 109 | captioned_image = caption_image_overlay(image, title, footer_left, footer_mid, footer_right) 110 | captioned_image = insert_image_data_embed(captioned_image, data) 111 | 112 | captioned_image.save(os.path.join(destination_folder, embedding_token+'.png'), "PNG", pnginfo=info) 113 | 114 | processed.images += [captioned_image] 115 | 116 | return processed 117 | --------------------------------------------------------------------------------