OPENGIT/stable-diffusion-webui
1
0
Fork 0
mirror of https://github.com/AUTOMATIC1111/stable-diffusion-webui.git synced 2025-02-07 06:02:53 +08:00
Code Issues Packages Projects Releases Wiki Activity

formatting

Browse Source
This commit is contained in:
DepFA 2022-10-12 13:13:25 +01:00 committed by GitHub
parent 66ec505975
commit 50be33e953
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 91 additions and 79 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
modules/textual_inversion/image_embedding.py
View File

@ -6,38 +6,46 @@ from PIL import Image,PngImagePlugin,ImageDraw,ImageFont
from fonts.ttf import Roboto from fonts.ttf import Roboto
import torch import torch
class EmbeddingEncoder(json.JSONEncoder): class EmbeddingEncoder(json.JSONEncoder):
def default(self, obj): def default(self, obj):
if isinstance(obj, torch.Tensor): if isinstance(obj, torch.Tensor):
return {'TORCHTENSOR': obj.cpu().detach().numpy().tolist()} return {'TORCHTENSOR': obj.cpu().detach().numpy().tolist()}
return json.JSONEncoder.default(self, obj) return json.JSONEncoder.default(self, obj)
class EmbeddingDecoder(json.JSONDecoder): class EmbeddingDecoder(json.JSONDecoder):
def __init__(self, *args, **kwargs): def __init__(self, *args, **kwargs):
json.JSONDecoder.__init__(self, object_hook=self.object_hook, *args, **kwargs) json.JSONDecoder.__init__(self, object_hook=self.object_hook, *args, **kwargs)
def object_hook(self, d): def object_hook(self, d):
if 'TORCHTENSOR' in d: if 'TORCHTENSOR' in d:
return torch.from_numpy(np.array(d['TORCHTENSOR'])) return torch.from_numpy(np.array(d['TORCHTENSOR']))
return d return d
def embedding_to_b64(data): def embedding_to_b64(data):
d = json.dumps(data, cls=EmbeddingEncoder) d = json.dumps(data, cls=EmbeddingEncoder)
return base64.b64encode(d.encode()) return base64.b64encode(d.encode())
def embedding_from_b64(data): def embedding_from_b64(data):
d = base64.b64decode(data) d = base64.b64decode(data)
return json.loads(d, cls=EmbeddingDecoder) return json.loads(d, cls=EmbeddingDecoder)
def lcg(m=2**32, a=1664525, c=1013904223, seed=0): def lcg(m=2**32, a=1664525, c=1013904223, seed=0):
while True: while True:
seed = (a * seed + c) % m seed = (a * seed + c) % m
yield seed % 255 yield seed % 255
def xor_block(block): def xor_block(block):
g = lcg() g = lcg()
randblock = np.array([next(g) for _ in range(np.product(block.shape))]).astype(np.uint8).reshape(block.shape) 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) return np.bitwise_xor(block.astype(np.uint8), randblock & 0x0F)
def style_block(block, sequence): def style_block(block, sequence):
im = Image.new('RGB', (block.shape[1], block.shape[0])) im = Image.new('RGB', (block.shape[1], block.shape[0]))
draw = ImageDraw.Draw(im) draw = ImageDraw.Draw(im)
@ -55,6 +63,7 @@ def style_block(block,sequence):
return block ^ fg return block ^ fg
def insert_image_data_embed(image, data): def insert_image_data_embed(image, data):
d = 3 d = 3
data_compressed = zlib.compress(json.dumps(data, cls=EmbeddingEncoder).encode(), level=9) data_compressed = zlib.compress(json.dumps(data, cls=EmbeddingEncoder).encode(), level=9)
@ -90,6 +99,7 @@ def insert_image_data_embed(image,data):
return background return background
def crop_black(img, tol=0): def crop_black(img, tol=0):
mask = (img > tol).all(2) mask = (img > tol).all(2)
mask0, mask1 = mask.any(0), mask.any(1) mask0, mask1 = mask.any(0), mask.any(1)
@ -97,6 +107,7 @@ def crop_black(img,tol=0):
row_start, row_end = mask1.argmax(), mask.shape[0]-mask1[::-1].argmax() row_start, row_end = mask1.argmax(), mask.shape[0]-mask1[::-1].argmax()
return img[row_start:row_end, col_start:col_end] return img[row_start:row_end, col_start:col_end]
def extract_image_data_embed(image): def extract_image_data_embed(image):
d = 3 d = 3
outarr = crop_black(np.array(image.convert('RGB').getdata()).reshape(image.size[1], image.size[0], d).astype(np.uint8)) & 0x0F outarr = crop_black(np.array(image.convert('RGB').getdata()).reshape(image.size[1], image.size[0], d).astype(np.uint8)) & 0x0F
@ -117,6 +128,7 @@ def extract_image_data_embed(image):
data = zlib.decompress(data_block) data = zlib.decompress(data_block)
return json.loads(data, cls=EmbeddingDecoder) return json.loads(data, cls=EmbeddingDecoder)
def caption_image_overlay(srcimage, title, footerLeft, footerMid, footerRight, textfont=None): def caption_image_overlay(srcimage, title, footerLeft, footerMid, footerRight, textfont=None):
from math import cos from math import cos
@ -163,10 +175,10 @@ def caption_image_overlay(srcimage,title,footerLeft,footerMid,footerRight,textfo
return image return image
if __name__ == '__main__': if __name__ == '__main__':
testEmbed = Image.open('test_embedding.png') testEmbed = Image.open('test_embedding.png')
data = extract_image_data_embed(testEmbed) data = extract_image_data_embed(testEmbed)
assert data is not None assert data is not None