mirror of
https://github.com/vladmandic/automatic
synced 2026-08-29 16:41:01 +02:00
Aptronymist
605988f520
Took out the print statement 'no image data blocks found' as it happens frequently as embeddings rarely have image data.
214 lines
8.1 KiB
Python
214 lines
8.1 KiB
Python
import base64
|
|
import json
|
|
import numpy as np
|
|
import zlib
|
|
from PIL import Image, PngImagePlugin, ImageDraw, ImageFont
|
|
import torch
|
|
from modules.shared import opts
|
|
|
|
|
|
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, *args, object_hook=self.object_hook, **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.prod(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 = np.resize(data_np_low, next_size)
|
|
data_np_low = data_np_low.reshape((h, -1, d))
|
|
|
|
data_np_high = np.resize(data_np_high, 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:
|
|
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()
|
|
fontsize = 32
|
|
if textfont is None:
|
|
textfont = opts.font or 'html/roboto.ttf'
|
|
|
|
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)
|
|
|
|
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
|