Spaces:
Running
on
Zero
Running
on
Zero
File size: 4,543 Bytes
96bea52 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 |
import os, json
import cv2
import glob
import numpy as np
from PIL import Image
import torch
def im_resize(original_image, short_len=1024):
h, w = original_image.shape[:-1]
if min(h, w) != short_len:
if h > w:
out_size = (short_len, int(h/w*short_len))
else:
out_size = (int(w/h*short_len), short_len)
else:
out_size = w, h
return cv2.resize(original_image, out_size)
def pixelize(image, block_size=64):
# 获取图像的宽度和高度
height, width, _ = image.shape
# 计算新图像的宽度和高度，使得每个块为 block_size x block_size 的大小
new_width = (width // block_size) * block_size
new_height = (height // block_size) * block_size
# 缩放图像以匹配新的宽度和高度
resized_image = cv2.resize(image, (new_width, new_height))
# 将图像分割成块并用块的平均值替代
for i in range(0, new_height, block_size):
for j in range(0, new_width, block_size):
block = resized_image[i:i+block_size, j:j+block_size, :]
average_color = np.mean(block, axis=(0, 1), dtype=int)
resized_image[i:i+block_size, j:j+block_size, :] = average_color
# 将图像缩小回原始大小，以增加像素风格的效果
final_image = cv2.resize(resized_image, (width, height))
return final_image
def get_kps_bbox_faceid(w, h, json_path):
def get_new_kps_and_bbox(w, h, kps, bbox):
scale = 512/max(w, h)
pad = abs(w - h) * scale / 2
if w < h:
kps[:, 0] -= pad
bbox[0] -= pad
bbox[2] -= pad
elif h < w:
kps[:, 1] -= pad
bbox[1] -= pad
bbox[3] -= pad
kps /= scale
bbox /= scale
return kps, bbox
with open(json_path, 'r') as file:
data = json.load(file)
kps = np.array(data.get("kps"))
bbox = np.array(data.get("bbox"))
kps, bbox = get_new_kps_and_bbox(w, h, kps, bbox)
embedding = data.get("embedding")
face_id_embed = embedding / np.linalg.norm(embedding)
face_id_embed = torch.from_numpy(face_id_embed)
return kps, bbox, face_id_embed
def get_kps_and_face_id_embed(w, h, json_path):
def get_new_kps(w, h, kps):
scale = 512/max(w, h)
pad = abs(w - h) * scale / 2
if w < h:
kps[:, 0] -= pad
elif h < w:
kps[:, 1] -= pad
kps = kps / scale
return kps
with open(json_path, 'r') as file:
data = json.load(file)
kps = np.array(data.get("kps"))
kps = get_new_kps(w, h, kps)
embedding = data.get("embedding")
face_id_embed = embedding / np.linalg.norm(embedding)
face_id_embed = torch.from_numpy(face_id_embed)
return kps, face_id_embed
def get_face_id_embed(json_path):
with open(json_path, 'r') as file:
data = json.load(file)
embedding = data.get("embedding")
face_id_embed = embedding / np.linalg.norm(embedding)
face_id_embed = torch.from_numpy(face_id_embed)
return face_id_embed
def get_module_kohya_state_dict(module, prefix: str, dtype: torch.dtype, adapter_name: str = "default"):
kohya_ss_state_dict = {}
for peft_key, weight in module.items():
kohya_key = peft_key.replace("unet.base_model.model", prefix)
kohya_key = kohya_key.replace("lora_A", "lora_down")
kohya_key = kohya_key.replace("lora_B", "lora_up")
kohya_key = kohya_key.replace(".", "_", kohya_key.count(".") - 2)
kohya_ss_state_dict[kohya_key] = weight.to(dtype)
# Set alpha parameter
if "lora_down" in kohya_key:
alpha_key = f'{kohya_key.split(".")[0]}.alpha'
kohya_ss_state_dict[alpha_key] = torch.tensor(8).to(dtype)
return kohya_ss_state_dict
def get_module_kohya_state_dict_xs(module, dtype):
kohya_ss_state_dict = {}
for peft_key, weight in module.items():
if "mid_block" in peft_key:
peft_key = peft_key.replace('attentions', 'base_midblock.attentions')
elif "down_block" in peft_key:
peft_key = peft_key.replace('attentions', 'base_attentions')
if dtype == None:
kohya_ss_state_dict[peft_key] = weight
else:
kohya_ss_state_dict[peft_key] = weight.to(dtype)
return kohya_ss_state_dict |