Merge branch 'master' into master

7 files changed, 930 insertions, 0 deletions

diff --git a/modules/textual_inversion/dataset.py b/modules/textual_inversion/dataset.py
new file mode 100644
index 00000000..23bb4b6a
--- /dev/null
+++ b/modules/textual_inversion/dataset.py
@@ -0,0 +1,121 @@
+import os

+import numpy as np

+import PIL

+import torch

+from PIL import Image

+from torch.utils.data import Dataset

+from torchvision import transforms

+

+import random

+import tqdm

+from modules import devices, shared

+import re

+

+re_numbers_at_start = re.compile(r"^[-\d]+\s*")

+

+

+class DatasetEntry:

+    def __init__(self, filename=None, latent=None, filename_text=None):

+        self.filename = filename

+        self.latent = latent

+        self.filename_text = filename_text

+        self.cond = None

+        self.cond_text = None

+

+

+class PersonalizedBase(Dataset):

+    def __init__(self, data_root, width, height, repeats, flip_p=0.5, placeholder_token="*", model=None, device=None, template_file=None, include_cond=False, batch_size=1):

+        re_word = re.compile(shared.opts.dataset_filename_word_regex) if len(shared.opts.dataset_filename_word_regex) > 0 else None

+

+        self.placeholder_token = placeholder_token

+

+        self.batch_size = batch_size

+        self.width = width

+        self.height = height

+        self.flip = transforms.RandomHorizontalFlip(p=flip_p)

+

+        self.dataset = []

+

+        with open(template_file, "r") as file:

+            lines = [x.strip() for x in file.readlines()]

+

+        self.lines = lines

+

+        assert data_root, 'dataset directory not specified'

+

+        cond_model = shared.sd_model.cond_stage_model

+

+        self.image_paths = [os.path.join(data_root, file_path) for file_path in os.listdir(data_root)]

+        print("Preparing dataset...")

+        for path in tqdm.tqdm(self.image_paths):

+            try:

+                image = Image.open(path).convert('RGB').resize((self.width, self.height), PIL.Image.BICUBIC)

+            except Exception:

+                continue

+

+            text_filename = os.path.splitext(path)[0] + ".txt"

+            filename = os.path.basename(path)

+

+            if os.path.exists(text_filename):

+                with open(text_filename, "r", encoding="utf8") as file:

+                    filename_text = file.read()

+            else:

+                filename_text = os.path.splitext(filename)[0]

+                filename_text = re.sub(re_numbers_at_start, '', filename_text)

+                if re_word:

+                    tokens = re_word.findall(filename_text)

+                    filename_text = (shared.opts.dataset_filename_join_string or "").join(tokens)

+

+            npimage = np.array(image).astype(np.uint8)

+            npimage = (npimage / 127.5 - 1.0).astype(np.float32)

+

+            torchdata = torch.from_numpy(npimage).to(device=device, dtype=torch.float32)

+            torchdata = torch.moveaxis(torchdata, 2, 0)

+

+            init_latent = model.get_first_stage_encoding(model.encode_first_stage(torchdata.unsqueeze(dim=0))).squeeze()

+            init_latent = init_latent.to(devices.cpu)

+

+            entry = DatasetEntry(filename=path, filename_text=filename_text, latent=init_latent)

+

+            if include_cond:

+                entry.cond_text = self.create_text(filename_text)

+                entry.cond = cond_model([entry.cond_text]).to(devices.cpu).squeeze(0)

+

+            self.dataset.append(entry)

+

+        assert len(self.dataset) > 1, "No images have been found in the dataset."

+        self.length = len(self.dataset) * repeats // batch_size

+

+        self.initial_indexes = np.arange(len(self.dataset))

+        self.indexes = None

+        self.shuffle()

+

+    def shuffle(self):

+        self.indexes = self.initial_indexes[torch.randperm(self.initial_indexes.shape[0])]

+

+    def create_text(self, filename_text):

+        text = random.choice(self.lines)

+        text = text.replace("[name]", self.placeholder_token)

+        text = text.replace("[filewords]", filename_text)

+        return text

+

+    def __len__(self):

+        return self.length

+

+    def __getitem__(self, i):

+        res = []

+

+        for j in range(self.batch_size):

+            position = i * self.batch_size + j

+            if position % len(self.indexes) == 0:

+                self.shuffle()

+

+            index = self.indexes[position % len(self.indexes)]

+            entry = self.dataset[index]

+

+            if entry.cond is None:

+                entry.cond_text = self.create_text(entry.filename_text)

+

+            res.append(entry)

+

+        return res

diff --git a/modules/textual_inversion/image_embedding.py b/modules/textual_inversion/image_embedding.py
new file mode 100644
index 00000000..898ce3b3
--- /dev/null
+++ b/modules/textual_inversion/image_embedding.py
@@ -0,0 +1,219 @@
+import base64

+import json

+import numpy as np

+import zlib

+from PIL import Image, PngImagePlugin, ImageDraw, ImageFont

+from fonts.ttf import Roboto

+import torch

+

+

+class EmbeddingEncoder(json.JSONEncoder):

+    def default(self, obj):

+        if isinstance(obj, torch.Tensor):

+            return {'TORCHTENSOR': obj.cpu().detach().numpy().tolist()}

+        return json.JSONEncoder.default(self, obj)

+

+

+class EmbeddingDecoder(json.JSONDecoder):

+    def __init__(self, *args, **kwargs):

+        json.JSONDecoder.__init__(self, object_hook=self.object_hook, *args, **kwargs)

+

+    def object_hook(self, d):

+        if 'TORCHTENSOR' in d:

+            return torch.from_numpy(np.array(d['TORCHTENSOR']))

+        return d

+

+

+def embedding_to_b64(data):

+    d = json.dumps(data, cls=EmbeddingEncoder)

+    return base64.b64encode(d.encode())

+

+

+def embedding_from_b64(data):

+    d = base64.b64decode(data)

+    return json.loads(d, cls=EmbeddingDecoder)

+

+

+def lcg(m=2**32, a=1664525, c=1013904223, seed=0):

+    while True:

+        seed = (a * seed + c) % m

+        yield seed % 255

+

+

+def xor_block(block):

+    g = lcg()

+    randblock = np.array([next(g) for _ in range(np.product(block.shape))]).astype(np.uint8).reshape(block.shape)

+    return np.bitwise_xor(block.astype(np.uint8), randblock & 0x0F)

+

+

+def style_block(block, sequence):

+    im = Image.new('RGB', (block.shape[1], block.shape[0]))

+    draw = ImageDraw.Draw(im)

+    i = 0

+    for x in range(-6, im.size[0], 8):

+        for yi, y in enumerate(range(-6, im.size[1], 8)):

+            offset = 0

+            if yi % 2 == 0:

+                offset = 4

+            shade = sequence[i % len(sequence)]

+            i += 1

+            draw.ellipse((x+offset, y, x+6+offset, y+6), fill=(shade, shade, shade))

+

+    fg = np.array(im).astype(np.uint8) & 0xF0

+

+    return block ^ fg

+

+

+def insert_image_data_embed(image, data):

+    d = 3

+    data_compressed = zlib.compress(json.dumps(data, cls=EmbeddingEncoder).encode(), level=9)

+    data_np_ = np.frombuffer(data_compressed, np.uint8).copy()

+    data_np_high = data_np_ >> 4

+    data_np_low = data_np_ & 0x0F

+

+    h = image.size[1]

+    next_size = data_np_low.shape[0] + (h-(data_np_low.shape[0] % h))

+    next_size = next_size + ((h*d)-(next_size % (h*d)))

+

+    data_np_low.resize(next_size)

+    data_np_low = data_np_low.reshape((h, -1, d))

+

+    data_np_high.resize(next_size)

+    data_np_high = data_np_high.reshape((h, -1, d))

+

+    edge_style = list(data['string_to_param'].values())[0].cpu().detach().numpy().tolist()[0][:1024]

+    edge_style = (np.abs(edge_style)/np.max(np.abs(edge_style))*255).astype(np.uint8)

+

+    data_np_low = style_block(data_np_low, sequence=edge_style)

+    data_np_low = xor_block(data_np_low)

+    data_np_high = style_block(data_np_high, sequence=edge_style[::-1])

+    data_np_high = xor_block(data_np_high)

+

+    im_low = Image.fromarray(data_np_low, mode='RGB')

+    im_high = Image.fromarray(data_np_high, mode='RGB')

+

+    background = Image.new('RGB', (image.size[0]+im_low.size[0]+im_high.size[0]+2, image.size[1]), (0, 0, 0))

+    background.paste(im_low, (0, 0))

+    background.paste(image, (im_low.size[0]+1, 0))

+    background.paste(im_high, (im_low.size[0]+1+image.size[0]+1, 0))

+

+    return background

+

+

+def crop_black(img, tol=0):

+    mask = (img > tol).all(2)

+    mask0, mask1 = mask.any(0), mask.any(1)

+    col_start, col_end = mask0.argmax(), mask.shape[1]-mask0[::-1].argmax()

+    row_start, row_end = mask1.argmax(), mask.shape[0]-mask1[::-1].argmax()

+    return img[row_start:row_end, col_start:col_end]

+

+

+def extract_image_data_embed(image):

+    d = 3

+    outarr = crop_black(np.array(image.convert('RGB').getdata()).reshape(image.size[1], image.size[0], d).astype(np.uint8)) & 0x0F

+    black_cols = np.where(np.sum(outarr, axis=(0, 2)) == 0)

+    if black_cols[0].shape[0] < 2:

+        print('No Image data blocks found.')

+        return None

+

+    data_block_lower = outarr[:, :black_cols[0].min(), :].astype(np.uint8)

+    data_block_upper = outarr[:, black_cols[0].max()+1:, :].astype(np.uint8)

+

+    data_block_lower = xor_block(data_block_lower)

+    data_block_upper = xor_block(data_block_upper)

+

+    data_block = (data_block_upper << 4) | (data_block_lower)

+    data_block = data_block.flatten().tobytes()

+

+    data = zlib.decompress(data_block)

+    return json.loads(data, cls=EmbeddingDecoder)

+

+

+def caption_image_overlay(srcimage, title, footerLeft, footerMid, footerRight, textfont=None):

+    from math import cos

+

+    image = srcimage.copy()

+

+    if textfont is None:

+        try:

+            textfont = ImageFont.truetype(opts.font or Roboto, fontsize)

+            textfont = opts.font or Roboto

+        except Exception:

+            textfont = Roboto

+

+    factor = 1.5

+    gradient = Image.new('RGBA', (1, image.size[1]), color=(0, 0, 0, 0))

+    for y in range(image.size[1]):

+        mag = 1-cos(y/image.size[1]*factor)

+        mag = max(mag, 1-cos((image.size[1]-y)/image.size[1]*factor*1.1))

+        gradient.putpixel((0, y), (0, 0, 0, int(mag*255)))

+    image = Image.alpha_composite(image.convert('RGBA'), gradient.resize(image.size))

+

+    draw = ImageDraw.Draw(image)

+    fontsize = 32

+    font = ImageFont.truetype(textfont, fontsize)

+    padding = 10

+

+    _, _, w, h = draw.textbbox((0, 0), title, font=font)

+    fontsize = min(int(fontsize * (((image.size[0]*0.75)-(padding*4))/w)), 72)

+    font = ImageFont.truetype(textfont, fontsize)

+    _, _, w, h = draw.textbbox((0, 0), title, font=font)

+    draw.text((padding, padding), title, anchor='lt', font=font, fill=(255, 255, 255, 230))

+

+    _, _, w, h = draw.textbbox((0, 0), footerLeft, font=font)

+    fontsize_left = min(int(fontsize * (((image.size[0]/3)-(padding))/w)), 72)

+    _, _, w, h = draw.textbbox((0, 0), footerMid, font=font)

+    fontsize_mid = min(int(fontsize * (((image.size[0]/3)-(padding))/w)), 72)

+    _, _, w, h = draw.textbbox((0, 0), footerRight, font=font)

+    fontsize_right = min(int(fontsize * (((image.size[0]/3)-(padding))/w)), 72)

+

+    font = ImageFont.truetype(textfont, min(fontsize_left, fontsize_mid, fontsize_right))

+

+    draw.text((padding, image.size[1]-padding),               footerLeft, anchor='ls', font=font, fill=(255, 255, 255, 230))

+    draw.text((image.size[0]/2, image.size[1]-padding),       footerMid, anchor='ms', font=font, fill=(255, 255, 255, 230))

+    draw.text((image.size[0]-padding, image.size[1]-padding), footerRight, anchor='rs', font=font, fill=(255, 255, 255, 230))

+

+    return image

+

+

+if __name__ == '__main__':

+

+    testEmbed = Image.open('test_embedding.png')

+    data = extract_image_data_embed(testEmbed)

+    assert data is not None

+

+    data = embedding_from_b64(testEmbed.text['sd-ti-embedding'])

+    assert data is not None

+

+    image = Image.new('RGBA', (512, 512), (255, 255, 200, 255))

+    cap_image = caption_image_overlay(image, 'title', 'footerLeft', 'footerMid', 'footerRight')

+

+    test_embed = {'string_to_param': {'*': torch.from_numpy(np.random.random((2, 4096)))}}

+

+    embedded_image = insert_image_data_embed(cap_image, test_embed)

+

+    retrived_embed = extract_image_data_embed(embedded_image)

+

+    assert str(retrived_embed) == str(test_embed)

+

+    embedded_image2 = insert_image_data_embed(cap_image, retrived_embed)

+

+    assert embedded_image == embedded_image2

+

+    g = lcg()

+    shared_random = np.array([next(g) for _ in range(100)]).astype(np.uint8).tolist()

+

+    reference_random = [253, 242, 127,  44, 157,  27, 239, 133,  38,  79, 167,   4, 177,

+                         95, 130,  79,  78,  14,  52, 215, 220, 194, 126,  28, 240, 179,

+                        160, 153, 149,  50, 105,  14,  21, 218, 199,  18,  54, 198, 193,

+                         38, 128,  19,  53, 195, 124,  75, 205,  12,   6, 145,   0,  28,

+                         30, 148,   8,  45, 218, 171,  55, 249,  97, 166,  12,  35,   0,

+                         41, 221, 122, 215, 170,  31, 113, 186,  97, 119,  31,  23, 185,

+                         66, 140,  30,  41,  37,  63, 137, 109, 216,  55, 159, 145,  82,

+                         204, 86,  73, 222,  44, 198, 118, 240,  97]

+

+    assert shared_random == reference_random

+

+    hunna_kay_random_sum = sum(np.array([next(g) for _ in range(100000)]).astype(np.uint8).tolist())

+

+    assert 12731374 == hunna_kay_random_sum

diff --git a/modules/textual_inversion/learn_schedule.py b/modules/textual_inversion/learn_schedule.py
new file mode 100644
index 00000000..2062726a
--- /dev/null
+++ b/modules/textual_inversion/learn_schedule.py
@@ -0,0 +1,69 @@
+import tqdm

+

+

+class LearnScheduleIterator:

+    def __init__(self, learn_rate, max_steps, cur_step=0):

+        """

+        specify learn_rate as "0.001:100, 0.00001:1000, 1e-5:10000" to have lr of 0.001 until step 100, 0.00001 until 1000, 1e-5:10000 until 10000

+        """

+

+        pairs = learn_rate.split(',')

+        self.rates = []

+        self.it = 0

+        self.maxit = 0

+        for i, pair in enumerate(pairs):

+            tmp = pair.split(':')

+            if len(tmp) == 2:

+                step = int(tmp[1])

+                if step > cur_step:

+                    self.rates.append((float(tmp[0]), min(step, max_steps)))

+                    self.maxit += 1

+                    if step > max_steps:

+                        return

+                elif step == -1:

+                    self.rates.append((float(tmp[0]), max_steps))

+                    self.maxit += 1

+                    return

+            else:

+                self.rates.append((float(tmp[0]), max_steps))

+                self.maxit += 1

+                return

+

+    def __iter__(self):

+        return self

+

+    def __next__(self):

+        if self.it < self.maxit:

+            self.it += 1

+            return self.rates[self.it - 1]

+        else:

+            raise StopIteration

+

+

+class LearnRateScheduler:

+    def __init__(self, learn_rate, max_steps, cur_step=0, verbose=True):

+        self.schedules = LearnScheduleIterator(learn_rate, max_steps, cur_step)

+        (self.learn_rate,  self.end_step) = next(self.schedules)

+        self.verbose = verbose

+

+        if self.verbose:

+            print(f'Training at rate of {self.learn_rate} until step {self.end_step}')

+

+        self.finished = False

+

+    def apply(self, optimizer, step_number):

+        if step_number <= self.end_step:

+            return

+

+        try:

+            (self.learn_rate, self.end_step) = next(self.schedules)

+        except Exception:

+            self.finished = True

+            return

+

+        if self.verbose:

+            tqdm.tqdm.write(f'Training at rate of {self.learn_rate} until step {self.end_step}')

+

+        for pg in optimizer.param_groups:

+            pg['lr'] = self.learn_rate

+

diff --git a/modules/textual_inversion/preprocess.py b/modules/textual_inversion/preprocess.py
new file mode 100644
index 00000000..886cf0c3
--- /dev/null
+++ b/modules/textual_inversion/preprocess.py
@@ -0,0 +1,116 @@
+import os

+from PIL import Image, ImageOps

+import platform

+import sys

+import tqdm

+import time

+

+from modules import shared, images

+from modules.shared import opts, cmd_opts

+if cmd_opts.deepdanbooru:

+    import modules.deepbooru as deepbooru

+

+

+def preprocess(process_src, process_dst, process_width, process_height, process_flip, process_split, process_caption, process_caption_deepbooru=False):

+    try:

+        if process_caption:

+            shared.interrogator.load()

+

+        if process_caption_deepbooru:

+            db_opts = deepbooru.create_deepbooru_opts()

+            db_opts[deepbooru.OPT_INCLUDE_RANKS] = False

+            deepbooru.create_deepbooru_process(opts.interrogate_deepbooru_score_threshold, db_opts)

+

+        preprocess_work(process_src, process_dst, process_width, process_height, process_flip, process_split, process_caption, process_caption_deepbooru)

+

+    finally:

+

+        if process_caption:

+            shared.interrogator.send_blip_to_ram()

+

+        if process_caption_deepbooru:

+            deepbooru.release_process()

+

+

+

+def preprocess_work(process_src, process_dst, process_width, process_height, process_flip, process_split, process_caption, process_caption_deepbooru=False):

+    width = process_width

+    height = process_height

+    src = os.path.abspath(process_src)

+    dst = os.path.abspath(process_dst)

+

+    assert src != dst, 'same directory specified as source and destination'

+

+    os.makedirs(dst, exist_ok=True)

+

+    files = os.listdir(src)

+

+    shared.state.textinfo = "Preprocessing..."

+    shared.state.job_count = len(files)

+

+    def save_pic_with_caption(image, index):

+        caption = ""

+

+        if process_caption:

+            caption += shared.interrogator.generate_caption(image)

+

+        if process_caption_deepbooru:

+            if len(caption) > 0:

+                caption += ", "

+            caption += deepbooru.get_tags_from_process(image)

+

+        filename_part = filename

+        filename_part = os.path.splitext(filename_part)[0]

+        filename_part = os.path.basename(filename_part)

+

+        basename = f"{index:05}-{subindex[0]}-{filename_part}"

+        image.save(os.path.join(dst, f"{basename}.png"))

+

+        if len(caption) > 0:

+            with open(os.path.join(dst, f"{basename}.txt"), "w", encoding="utf8") as file:

+                file.write(caption)

+

+        subindex[0] += 1

+

+    def save_pic(image, index):

+        save_pic_with_caption(image, index)

+

+        if process_flip:

+            save_pic_with_caption(ImageOps.mirror(image), index)

+

+    for index, imagefile in enumerate(tqdm.tqdm(files)):

+        subindex = [0]

+        filename = os.path.join(src, imagefile)

+        try:

+            img = Image.open(filename).convert("RGB")

+        except Exception:

+            continue

+

+        if shared.state.interrupted:

+            break

+

+        ratio = img.height / img.width

+        is_tall = ratio > 1.35

+        is_wide = ratio < 1 / 1.35

+

+        if process_split and is_tall:

+            img = img.resize((width, height * img.height // img.width))

+

+            top = img.crop((0, 0, width, height))

+            save_pic(top, index)

+

+            bot = img.crop((0, img.height - height, width, img.height))

+            save_pic(bot, index)

+        elif process_split and is_wide:

+            img = img.resize((width * img.width // img.height, height))

+

+            left = img.crop((0, 0, width, height))

+            save_pic(left, index)

+

+            right = img.crop((img.width - width, 0, img.width, height))

+            save_pic(right, index)

+        else:

+            img = images.resize_image(1, img, width, height)

+            save_pic(img, index)

+

+        shared.state.nextjob()

diff --git a/modules/textual_inversion/test_embedding.png b/modules/textual_inversion/test_embedding.png
new file mode 100644
index 00000000..07e2d9af
--- /dev/null
+++ b/modules/textual_inversion/test_embedding.png
diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py
new file mode 100644
index 00000000..e754747e
--- /dev/null
+++ b/modules/textual_inversion/textual_inversion.py
@@ -0,0 +1,363 @@
+import os

+import sys

+import traceback

+

+import torch

+import tqdm

+import html

+import datetime

+import csv

+

+from PIL import Image, PngImagePlugin

+

+from modules import shared, devices, sd_hijack, processing, sd_models

+import modules.textual_inversion.dataset

+from modules.textual_inversion.learn_schedule import LearnRateScheduler

+

+from modules.textual_inversion.image_embedding import (embedding_to_b64, embedding_from_b64,

+                                                       insert_image_data_embed, extract_image_data_embed,

+                                                       caption_image_overlay)

+

+class Embedding:

+    def __init__(self, vec, name, step=None):

+        self.vec = vec

+        self.name = name

+        self.step = step

+        self.cached_checksum = None

+        self.sd_checkpoint = None

+        self.sd_checkpoint_name = None

+

+    def save(self, filename):

+        embedding_data = {

+            "string_to_token": {"*": 265},

+            "string_to_param": {"*": self.vec},

+            "name": self.name,

+            "step": self.step,

+            "sd_checkpoint": self.sd_checkpoint,

+            "sd_checkpoint_name": self.sd_checkpoint_name,

+        }

+

+        torch.save(embedding_data, filename)

+

+    def checksum(self):

+        if self.cached_checksum is not None:

+            return self.cached_checksum

+

+        def const_hash(a):

+            r = 0

+            for v in a:

+                r = (r * 281 ^ int(v) * 997) & 0xFFFFFFFF

+            return r

+

+        self.cached_checksum = f'{const_hash(self.vec.reshape(-1) * 100) & 0xffff:04x}'

+        return self.cached_checksum

+

+

+class EmbeddingDatabase:

+    def __init__(self, embeddings_dir):

+        self.ids_lookup = {}

+        self.word_embeddings = {}

+        self.dir_mtime = None

+        self.embeddings_dir = embeddings_dir

+

+    def register_embedding(self, embedding, model):

+

+        self.word_embeddings[embedding.name] = embedding

+

+        ids = model.cond_stage_model.tokenizer([embedding.name], add_special_tokens=False)['input_ids'][0]

+

+        first_id = ids[0]

+        if first_id not in self.ids_lookup:

+            self.ids_lookup[first_id] = []

+

+        self.ids_lookup[first_id] = sorted(self.ids_lookup[first_id] + [(ids, embedding)], key=lambda x: len(x[0]), reverse=True)

+

+        return embedding

+

+    def load_textual_inversion_embeddings(self):

+        mt = os.path.getmtime(self.embeddings_dir)

+        if self.dir_mtime is not None and mt <= self.dir_mtime:

+            return

+

+        self.dir_mtime = mt

+        self.ids_lookup.clear()

+        self.word_embeddings.clear()

+

+        def process_file(path, filename):

+            name = os.path.splitext(filename)[0]

+

+            data = []

+

+            if filename.upper().endswith('.PNG'):

+                embed_image = Image.open(path)

+                if 'sd-ti-embedding' in embed_image.text:

+                    data = embedding_from_b64(embed_image.text['sd-ti-embedding'])

+                    name = data.get('name', name)

+                else:

+                    data = extract_image_data_embed(embed_image)

+                    name = data.get('name', name)

+            else:

+                data = torch.load(path, map_location="cpu")

+

+            # textual inversion embeddings

+            if 'string_to_param' in data:

+                param_dict = data['string_to_param']

+                if hasattr(param_dict, '_parameters'):

+                    param_dict = getattr(param_dict, '_parameters')  # fix for torch 1.12.1 loading saved file from torch 1.11

+                assert len(param_dict) == 1, 'embedding file has multiple terms in it'

+                emb = next(iter(param_dict.items()))[1]

+            # diffuser concepts

+            elif type(data) == dict and type(next(iter(data.values()))) == torch.Tensor:

+                assert len(data.keys()) == 1, 'embedding file has multiple terms in it'

+

+                emb = next(iter(data.values()))

+                if len(emb.shape) == 1:

+                    emb = emb.unsqueeze(0)

+            else:

+                raise Exception(f"Couldn't identify {filename} as neither textual inversion embedding nor diffuser concept.")

+

+            vec = emb.detach().to(devices.device, dtype=torch.float32)

+            embedding = Embedding(vec, name)

+            embedding.step = data.get('step', None)

+            embedding.sd_checkpoint = data.get('hash', None)

+            embedding.sd_checkpoint_name = data.get('sd_checkpoint_name', None)

+            self.register_embedding(embedding, shared.sd_model)

+

+        for fn in os.listdir(self.embeddings_dir):

+            try:

+                fullfn = os.path.join(self.embeddings_dir, fn)

+

+                if os.stat(fullfn).st_size == 0:

+                    continue

+

+                process_file(fullfn, fn)

+            except Exception:

+                print(f"Error loading emedding {fn}:", file=sys.stderr)

+                print(traceback.format_exc(), file=sys.stderr)

+                continue

+

+        print(f"Loaded a total of {len(self.word_embeddings)} textual inversion embeddings.")

+

+    def find_embedding_at_position(self, tokens, offset):

+        token = tokens[offset]

+        possible_matches = self.ids_lookup.get(token, None)

+

+        if possible_matches is None:

+            return None, None

+

+        for ids, embedding in possible_matches:

+            if tokens[offset:offset + len(ids)] == ids:

+                return embedding, len(ids)

+

+        return None, None

+

+

+def create_embedding(name, num_vectors_per_token, init_text='*'):

+    cond_model = shared.sd_model.cond_stage_model

+    embedding_layer = cond_model.wrapped.transformer.text_model.embeddings

+

+    ids = cond_model.tokenizer(init_text, max_length=num_vectors_per_token, return_tensors="pt", add_special_tokens=False)["input_ids"]

+    embedded = embedding_layer.token_embedding.wrapped(ids.to(devices.device)).squeeze(0)

+    vec = torch.zeros((num_vectors_per_token, embedded.shape[1]), device=devices.device)

+

+    for i in range(num_vectors_per_token):

+        vec[i] = embedded[i * int(embedded.shape[0]) // num_vectors_per_token]

+

+    fn = os.path.join(shared.cmd_opts.embeddings_dir, f"{name}.pt")

+    assert not os.path.exists(fn), f"file {fn} already exists"

+

+    embedding = Embedding(vec, name)

+    embedding.step = 0

+    embedding.save(fn)

+

+    return fn

+

+

+def write_loss(log_directory, filename, step, epoch_len, values):

+    if shared.opts.training_write_csv_every == 0:

+        return

+

+    if step % shared.opts.training_write_csv_every != 0:

+        return

+

+    write_csv_header = False if os.path.exists(os.path.join(log_directory, filename)) else True

+

+    with open(os.path.join(log_directory, filename), "a+", newline='') as fout:

+        csv_writer = csv.DictWriter(fout, fieldnames=["step", "epoch", "epoch_step", *(values.keys())])

+

+        if write_csv_header:

+            csv_writer.writeheader()

+

+        epoch = step // epoch_len

+        epoch_step = step - epoch * epoch_len

+

+        csv_writer.writerow({

+            "step": step + 1,

+            "epoch": epoch + 1,

+            "epoch_step": epoch_step + 1,

+            **values,

+        })

+

+

+def train_embedding(embedding_name, learn_rate, batch_size, data_root, log_directory, training_width, training_height, steps, create_image_every, save_embedding_every, template_file, save_image_with_stored_embedding, preview_from_txt2img, preview_prompt, preview_negative_prompt, preview_steps, preview_sampler_index, preview_cfg_scale, preview_seed, preview_width, preview_height):

+    assert embedding_name, 'embedding not selected'

+

+    shared.state.textinfo = "Initializing textual inversion training..."

+    shared.state.job_count = steps

+

+    filename = os.path.join(shared.cmd_opts.embeddings_dir, f'{embedding_name}.pt')

+

+    log_directory = os.path.join(log_directory, datetime.datetime.now().strftime("%Y-%m-%d"), embedding_name)

+

+    if save_embedding_every > 0:

+        embedding_dir = os.path.join(log_directory, "embeddings")

+        os.makedirs(embedding_dir, exist_ok=True)

+    else:

+        embedding_dir = None

+

+    if create_image_every > 0:

+        images_dir = os.path.join(log_directory, "images")

+        os.makedirs(images_dir, exist_ok=True)

+    else:

+        images_dir = None

+

+    if create_image_every > 0 and save_image_with_stored_embedding:

+        images_embeds_dir = os.path.join(log_directory, "image_embeddings")

+        os.makedirs(images_embeds_dir, exist_ok=True)

+    else:

+        images_embeds_dir = None

+        

+    cond_model = shared.sd_model.cond_stage_model

+

+    shared.state.textinfo = f"Preparing dataset from {html.escape(data_root)}..."

+    with torch.autocast("cuda"):

+        ds = modules.textual_inversion.dataset.PersonalizedBase(data_root=data_root, width=training_width, height=training_height, repeats=shared.opts.training_image_repeats_per_epoch, placeholder_token=embedding_name, model=shared.sd_model, device=devices.device, template_file=template_file, batch_size=batch_size)

+

+    hijack = sd_hijack.model_hijack

+

+    embedding = hijack.embedding_db.word_embeddings[embedding_name]

+    embedding.vec.requires_grad = True

+

+    losses = torch.zeros((32,))

+

+    last_saved_file = "<none>"

+    last_saved_image = "<none>"

+

+    ititial_step = embedding.step or 0

+    if ititial_step > steps:

+        return embedding, filename

+

+    scheduler = LearnRateScheduler(learn_rate, steps, ititial_step)

+    optimizer = torch.optim.AdamW([embedding.vec], lr=scheduler.learn_rate)

+

+    pbar = tqdm.tqdm(enumerate(ds), total=steps-ititial_step)

+    for i, entries in pbar:

+        embedding.step = i + ititial_step

+

+        scheduler.apply(optimizer, embedding.step)

+        if scheduler.finished:

+            break

+

+        if shared.state.interrupted:

+            break

+

+        with torch.autocast("cuda"):

+            c = cond_model([entry.cond_text for entry in entries])

+            x = torch.stack([entry.latent for entry in entries]).to(devices.device)

+            loss = shared.sd_model(x, c)[0]

+            del x

+

+            losses[embedding.step % losses.shape[0]] = loss.item()

+

+            optimizer.zero_grad()

+            loss.backward()

+            optimizer.step()

+

+        epoch_num = embedding.step // len(ds)

+        epoch_step = embedding.step - (epoch_num * len(ds)) + 1

+

+        pbar.set_description(f"[Epoch {epoch_num}: {epoch_step}/{len(ds)}]loss: {losses.mean():.7f}")

+

+        if embedding.step > 0 and embedding_dir is not None and embedding.step % save_embedding_every == 0:

+            last_saved_file = os.path.join(embedding_dir, f'{embedding_name}-{embedding.step}.pt')

+            embedding.save(last_saved_file)

+

+        write_loss(log_directory, "textual_inversion_loss.csv", embedding.step, len(ds), {

+            "loss": f"{losses.mean():.7f}",

+            "learn_rate": scheduler.learn_rate

+        })

+

+        if embedding.step > 0 and images_dir is not None and embedding.step % create_image_every == 0:

+            last_saved_image = os.path.join(images_dir, f'{embedding_name}-{embedding.step}.png')

+

+            p = processing.StableDiffusionProcessingTxt2Img(

+                sd_model=shared.sd_model,

+                do_not_save_grid=True,

+                do_not_save_samples=True,

+            )

+

+            if preview_from_txt2img:

+                p.prompt = preview_prompt

+                p.negative_prompt = preview_negative_prompt

+                p.steps = preview_steps

+                p.sampler_index = preview_sampler_index

+                p.cfg_scale = preview_cfg_scale

+                p.seed = preview_seed

+                p.width = preview_width

+                p.height = preview_height

+            else:

+                p.prompt = entries[0].cond_text

+                p.steps = 20

+                p.width = training_width

+                p.height = training_height

+

+            preview_text = p.prompt

+

+            processed = processing.process_images(p)

+            image = processed.images[0]

+

+            shared.state.current_image = image

+

+            if save_image_with_stored_embedding and os.path.exists(last_saved_file):

+

+                last_saved_image_chunks = os.path.join(images_embeds_dir, f'{embedding_name}-{embedding.step}.png')

+

+                info = PngImagePlugin.PngInfo()

+                data = torch.load(last_saved_file)

+                info.add_text("sd-ti-embedding", embedding_to_b64(data))

+

+                title = "<{}>".format(data.get('name', '???'))

+                checkpoint = sd_models.select_checkpoint()

+                footer_left = checkpoint.model_name

+                footer_mid = '[{}]'.format(checkpoint.hash)

+                footer_right = '{}'.format(embedding.step)

+

+                captioned_image = caption_image_overlay(image, title, footer_left, footer_mid, footer_right)

+                captioned_image = insert_image_data_embed(captioned_image, data)

+

+                captioned_image.save(last_saved_image_chunks, "PNG", pnginfo=info)

+

+            image.save(last_saved_image)

+

+            last_saved_image += f", prompt: {preview_text}"

+

+        shared.state.job_no = embedding.step

+

+        shared.state.textinfo = f"""

+<p>

+Loss: {losses.mean():.7f}<br/>

+Step: {embedding.step}<br/>

+Last prompt: {html.escape(entries[0].cond_text)}<br/>

+Last saved embedding: {html.escape(last_saved_file)}<br/>

+Last saved image: {html.escape(last_saved_image)}<br/>

+</p>

+"""

+

+    checkpoint = sd_models.select_checkpoint()

+

+    embedding.sd_checkpoint = checkpoint.hash

+    embedding.sd_checkpoint_name = checkpoint.model_name

+    embedding.cached_checksum = None

+    embedding.save(filename)

+

+    return embedding, filename

diff --git a/modules/textual_inversion/ui.py b/modules/textual_inversion/ui.py
new file mode 100644
index 00000000..36881e7a
--- /dev/null
+++ b/modules/textual_inversion/ui.py
@@ -0,0 +1,42 @@
+import html

+

+import gradio as gr

+

+import modules.textual_inversion.textual_inversion

+import modules.textual_inversion.preprocess

+from modules import sd_hijack, shared

+

+

+def create_embedding(name, initialization_text, nvpt):

+    filename = modules.textual_inversion.textual_inversion.create_embedding(name, nvpt, init_text=initialization_text)

+

+    sd_hijack.model_hijack.embedding_db.load_textual_inversion_embeddings()

+

+    return gr.Dropdown.update(choices=sorted(sd_hijack.model_hijack.embedding_db.word_embeddings.keys())), f"Created: {filename}", ""

+

+

+def preprocess(*args):

+    modules.textual_inversion.preprocess.preprocess(*args)

+

+    return "Preprocessing finished.", ""

+

+

+def train_embedding(*args):

+

+    assert not shared.cmd_opts.lowvram, 'Training models with lowvram not possible'

+

+    try:

+        sd_hijack.undo_optimizations()

+

+        embedding, filename = modules.textual_inversion.textual_inversion.train_embedding(*args)

+

+        res = f"""

+Training {'interrupted' if shared.state.interrupted else 'finished'} at {embedding.step} steps.

+Embedding saved to {html.escape(filename)}

+"""

+        return res, ""

+    except Exception:

+        raise

+    finally:

+        sd_hijack.apply_optimizations()

+