shift embedding logic out of textual_inversion

1 files changed, 6 insertions, 119 deletions

diff --git a/modules/textual_inversion/textual_inversion.py b/modules/textual_inversion/textual_inversion.py
index 8c66aeb5..22b4ae7f 100644
--- a/modules/textual_inversion/textual_inversion.py
+++ b/modules/textual_inversion/textual_inversion.py
@@ -7,124 +7,11 @@ import tqdm
 import html

 import datetime

 

-from PIL import Image,PngImagePlugin,ImageDraw

-from ..images import captionImageOverlay

-import numpy as np

-import base64

-import json

-import zlib

+from PIL import Image,PngImagePlugin

 

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

 import modules.textual_inversion.dataset

 

-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 embeddingToB64(data):

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

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

-

-def embeddingFromB64(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

-

-def xorBlock(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 styleBlock(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 insertImageDataEmbed(image,data):

-    d = 3

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

-    dnp = np.frombuffer(data_compressed,np.uint8).copy()

-    dnphigh = dnp >> 4

-    dnplow  = dnp & 0x0F

-    

-    h = image.size[1]

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

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

-

-    dnplow.resize(next_size)

-    dnplow = dnplow.reshape((h,-1,d))

-

-    dnphigh.resize(next_size)

-    dnphigh = dnphigh.reshape((h,-1,d))

-

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

-    edgeStyleWeights = (np.abs(edgeStyleWeights)/np.max(np.abs(edgeStyleWeights))*255).astype(np.uint8)

-

-    dnplow   = styleBlock(dnplow,sequence=edgeStyleWeights)

-    dnplow   = xorBlock(dnplow)    

-    dnphigh  = styleBlock(dnphigh,sequence=edgeStyleWeights[::-1])

-    dnphigh  = xorBlock(dnphigh)

-

-    imlow  = Image.fromarray(dnplow,mode='RGB')

-    imhigh = Image.fromarray(dnphigh,mode='RGB')

-

-    background = Image.new('RGB',(image.size[0]+imlow.size[0]+imhigh.size[0]+2,image.size[1]),(0,0,0))

-    background.paste(imlow,(0,0))

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

-    background.paste(imhigh,(imlow.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 extractImageDataEmbed(image):

-    d=3

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

-    blackCols = np.where( np.sum(outarr, axis=(0,2))==0)

-    if blackCols[0].shape[0] < 2:

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

-        return None

-

-    dataBlocklower = outarr[:,:blackCols[0].min(),:].astype(np.uint8)

-    dataBlockupper = outarr[:,blackCols[0].max()+1:,:].astype(np.uint8)

-

-    dataBlocklower = xorBlock(dataBlocklower)

-    dataBlockupper = xorBlock(dataBlockupper)

-    

-    dataBlock = (dataBlockupper << 4) | (dataBlocklower)

-    dataBlock = dataBlock.flatten().tobytes()   

-    data = zlib.decompress(dataBlock)

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

 

 class Embedding:

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

@@ -199,10 +86,10 @@ class EmbeddingDatabase:
             if filename.upper().endswith('.PNG'):

                 embed_image = Image.open(path)

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

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

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

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

                 else:

-                    data = extractImageDataEmbed(embed_image)

+                    data = extract_image_data_embed(embed_image)

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

             else:

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

@@ -393,7 +280,7 @@ def train_embedding(embedding_name, learn_rate, data_root, log_directory, traini
 

                 info = PngImagePlugin.PngInfo()

                 data = torch.load(last_saved_file)

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

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

 

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

                 checkpoint = sd_models.select_checkpoint()

@@ -401,8 +288,8 @@ def train_embedding(embedding_name, learn_rate, data_root, log_directory, traini
                 footer_mid = '[{}]'.format(checkpoint.hash)

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

 

-                captioned_image = captionImageOverlay(image,title,footer_left,footer_mid,footer_right)

-                captioned_image = insertImageDataEmbed(captioned_image,data)

+                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)