|
|
import os
|
|
|
import cv2
|
|
|
import json
|
|
|
import random
|
|
|
import glob
|
|
|
import torch
|
|
|
import einops
|
|
|
import numpy as np
|
|
|
import datetime
|
|
|
import torchvision
|
|
|
|
|
|
import safetensors.torch as sf
|
|
|
from PIL import Image
|
|
|
|
|
|
|
|
|
def min_resize(x, m):
|
|
|
if x.shape[0] < x.shape[1]:
|
|
|
s0 = m
|
|
|
s1 = int(float(m) / float(x.shape[0]) * float(x.shape[1]))
|
|
|
else:
|
|
|
s0 = int(float(m) / float(x.shape[1]) * float(x.shape[0]))
|
|
|
s1 = m
|
|
|
new_max = max(s1, s0)
|
|
|
raw_max = max(x.shape[0], x.shape[1])
|
|
|
if new_max < raw_max:
|
|
|
interpolation = cv2.INTER_AREA
|
|
|
else:
|
|
|
interpolation = cv2.INTER_LANCZOS4
|
|
|
y = cv2.resize(x, (s1, s0), interpolation=interpolation)
|
|
|
return y
|
|
|
|
|
|
|
|
|
def d_resize(x, y):
|
|
|
H, W, C = y.shape
|
|
|
new_min = min(H, W)
|
|
|
raw_min = min(x.shape[0], x.shape[1])
|
|
|
if new_min < raw_min:
|
|
|
interpolation = cv2.INTER_AREA
|
|
|
else:
|
|
|
interpolation = cv2.INTER_LANCZOS4
|
|
|
y = cv2.resize(x, (W, H), interpolation=interpolation)
|
|
|
return y
|
|
|
|
|
|
|
|
|
def resize_and_center_crop(image, target_width, target_height):
|
|
|
if target_height == image.shape[0] and target_width == image.shape[1]:
|
|
|
return image
|
|
|
|
|
|
pil_image = Image.fromarray(image)
|
|
|
original_width, original_height = pil_image.size
|
|
|
scale_factor = max(target_width / original_width, target_height / original_height)
|
|
|
resized_width = int(round(original_width * scale_factor))
|
|
|
resized_height = int(round(original_height * scale_factor))
|
|
|
resized_image = pil_image.resize((resized_width, resized_height), Image.LANCZOS)
|
|
|
left = (resized_width - target_width) / 2
|
|
|
top = (resized_height - target_height) / 2
|
|
|
right = (resized_width + target_width) / 2
|
|
|
bottom = (resized_height + target_height) / 2
|
|
|
cropped_image = resized_image.crop((left, top, right, bottom))
|
|
|
return np.array(cropped_image)
|
|
|
|
|
|
|
|
|
def resize_and_center_crop_pytorch(image, target_width, target_height):
|
|
|
B, C, H, W = image.shape
|
|
|
|
|
|
if H == target_height and W == target_width:
|
|
|
return image
|
|
|
|
|
|
scale_factor = max(target_width / W, target_height / H)
|
|
|
resized_width = int(round(W * scale_factor))
|
|
|
resized_height = int(round(H * scale_factor))
|
|
|
|
|
|
resized = torch.nn.functional.interpolate(image, size=(resized_height, resized_width), mode='bilinear', align_corners=False)
|
|
|
|
|
|
top = (resized_height - target_height) // 2
|
|
|
left = (resized_width - target_width) // 2
|
|
|
cropped = resized[:, :, top:top + target_height, left:left + target_width]
|
|
|
|
|
|
return cropped
|
|
|
|
|
|
|
|
|
def resize_without_crop(image, target_width, target_height):
|
|
|
if target_height == image.shape[0] and target_width == image.shape[1]:
|
|
|
return image
|
|
|
|
|
|
pil_image = Image.fromarray(image)
|
|
|
resized_image = pil_image.resize((target_width, target_height), Image.LANCZOS)
|
|
|
return np.array(resized_image)
|
|
|
|
|
|
|
|
|
def just_crop(image, w, h):
|
|
|
if h == image.shape[0] and w == image.shape[1]:
|
|
|
return image
|
|
|
|
|
|
original_height, original_width = image.shape[:2]
|
|
|
k = min(original_height / h, original_width / w)
|
|
|
new_width = int(round(w * k))
|
|
|
new_height = int(round(h * k))
|
|
|
x_start = (original_width - new_width) // 2
|
|
|
y_start = (original_height - new_height) // 2
|
|
|
cropped_image = image[y_start:y_start + new_height, x_start:x_start + new_width]
|
|
|
return cropped_image
|
|
|
|
|
|
|
|
|
def write_to_json(data, file_path):
|
|
|
temp_file_path = file_path + ".tmp"
|
|
|
with open(temp_file_path, 'wt', encoding='utf-8') as temp_file:
|
|
|
json.dump(data, temp_file, indent=4)
|
|
|
os.replace(temp_file_path, file_path)
|
|
|
return
|
|
|
|
|
|
|
|
|
def read_from_json(file_path):
|
|
|
with open(file_path, 'rt', encoding='utf-8') as file:
|
|
|
data = json.load(file)
|
|
|
return data
|
|
|
|
|
|
|
|
|
def get_active_parameters(m):
|
|
|
return {k: v for k, v in m.named_parameters() if v.requires_grad}
|
|
|
|
|
|
|
|
|
def cast_training_params(m, dtype=torch.float32):
|
|
|
result = {}
|
|
|
for n, param in m.named_parameters():
|
|
|
if param.requires_grad:
|
|
|
param.data = param.to(dtype)
|
|
|
result[n] = param
|
|
|
return result
|
|
|
|
|
|
|
|
|
def separate_lora_AB(parameters, B_patterns=None):
|
|
|
parameters_normal = {}
|
|
|
parameters_B = {}
|
|
|
|
|
|
if B_patterns is None:
|
|
|
B_patterns = ['.lora_B.', '__zero__']
|
|
|
|
|
|
for k, v in parameters.items():
|
|
|
if any(B_pattern in k for B_pattern in B_patterns):
|
|
|
parameters_B[k] = v
|
|
|
else:
|
|
|
parameters_normal[k] = v
|
|
|
|
|
|
return parameters_normal, parameters_B
|
|
|
|
|
|
|
|
|
def set_attr_recursive(obj, attr, value):
|
|
|
attrs = attr.split(".")
|
|
|
for name in attrs[:-1]:
|
|
|
obj = getattr(obj, name)
|
|
|
setattr(obj, attrs[-1], value)
|
|
|
return
|
|
|
|
|
|
|
|
|
def print_tensor_list_size(tensors):
|
|
|
total_size = 0
|
|
|
total_elements = 0
|
|
|
|
|
|
if isinstance(tensors, dict):
|
|
|
tensors = tensors.values()
|
|
|
|
|
|
for tensor in tensors:
|
|
|
total_size += tensor.nelement() * tensor.element_size()
|
|
|
total_elements += tensor.nelement()
|
|
|
|
|
|
total_size_MB = total_size / (1024 ** 2)
|
|
|
total_elements_B = total_elements / 1e9
|
|
|
|
|
|
print(f"Total number of tensors: {len(tensors)}")
|
|
|
print(f"Total size of tensors: {total_size_MB:.2f} MB")
|
|
|
print(f"Total number of parameters: {total_elements_B:.3f} billion")
|
|
|
return
|
|
|
|
|
|
|
|
|
@torch.no_grad()
|
|
|
def batch_mixture(a, b=None, probability_a=0.5, mask_a=None):
|
|
|
batch_size = a.size(0)
|
|
|
|
|
|
if b is None:
|
|
|
b = torch.zeros_like(a)
|
|
|
|
|
|
if mask_a is None:
|
|
|
mask_a = torch.rand(batch_size) < probability_a
|
|
|
|
|
|
mask_a = mask_a.to(a.device)
|
|
|
mask_a = mask_a.reshape((batch_size,) + (1,) * (a.dim() - 1))
|
|
|
result = torch.where(mask_a, a, b)
|
|
|
return result
|
|
|
|
|
|
|
|
|
@torch.no_grad()
|
|
|
def zero_module(module):
|
|
|
for p in module.parameters():
|
|
|
p.detach().zero_()
|
|
|
return module
|
|
|
|
|
|
|
|
|
@torch.no_grad()
|
|
|
def supress_lower_channels(m, k, alpha=0.01):
|
|
|
data = m.weight.data.clone()
|
|
|
|
|
|
assert int(data.shape[1]) >= k
|
|
|
|
|
|
data[:, :k] = data[:, :k] * alpha
|
|
|
m.weight.data = data.contiguous().clone()
|
|
|
return m
|
|
|
|
|
|
|
|
|
def freeze_module(m):
|
|
|
if not hasattr(m, '_forward_inside_frozen_module'):
|
|
|
m._forward_inside_frozen_module = m.forward
|
|
|
m.requires_grad_(False)
|
|
|
m.forward = torch.no_grad()(m.forward)
|
|
|
return m
|
|
|
|
|
|
|
|
|
def get_latest_safetensors(folder_path):
|
|
|
safetensors_files = glob.glob(os.path.join(folder_path, '*.safetensors'))
|
|
|
|
|
|
if not safetensors_files:
|
|
|
raise ValueError('No file to resume!')
|
|
|
|
|
|
latest_file = max(safetensors_files, key=os.path.getmtime)
|
|
|
latest_file = os.path.abspath(os.path.realpath(latest_file))
|
|
|
return latest_file
|
|
|
|
|
|
|
|
|
def generate_random_prompt_from_tags(tags_str, min_length=3, max_length=32):
|
|
|
tags = tags_str.split(', ')
|
|
|
tags = random.sample(tags, k=min(random.randint(min_length, max_length), len(tags)))
|
|
|
prompt = ', '.join(tags)
|
|
|
return prompt
|
|
|
|
|
|
|
|
|
def interpolate_numbers(a, b, n, round_to_int=False, gamma=1.0):
|
|
|
numbers = a + (b - a) * (np.linspace(0, 1, n) ** gamma)
|
|
|
if round_to_int:
|
|
|
numbers = np.round(numbers).astype(int)
|
|
|
return numbers.tolist()
|
|
|
|
|
|
|
|
|
def uniform_random_by_intervals(inclusive, exclusive, n, round_to_int=False):
|
|
|
edges = np.linspace(0, 1, n + 1)
|
|
|
points = np.random.uniform(edges[:-1], edges[1:])
|
|
|
numbers = inclusive + (exclusive - inclusive) * points
|
|
|
if round_to_int:
|
|
|
numbers = np.round(numbers).astype(int)
|
|
|
return numbers.tolist()
|
|
|
|
|
|
|
|
|
def soft_append_bcthw(history, current, overlap=0):
|
|
|
if overlap <= 0:
|
|
|
return torch.cat([history, current], dim=2)
|
|
|
|
|
|
assert history.shape[2] >= overlap, f"History length ({history.shape[2]}) must be >= overlap ({overlap})"
|
|
|
assert current.shape[2] >= overlap, f"Current length ({current.shape[2]}) must be >= overlap ({overlap})"
|
|
|
|
|
|
weights = torch.linspace(1, 0, overlap, dtype=history.dtype, device=history.device).view(1, 1, -1, 1, 1)
|
|
|
blended = weights * history[:, :, -overlap:] + (1 - weights) * current[:, :, :overlap]
|
|
|
output = torch.cat([history[:, :, :-overlap], blended, current[:, :, overlap:]], dim=2)
|
|
|
|
|
|
return output.to(history)
|
|
|
|
|
|
|
|
|
def save_bcthw_as_mp4(x, output_filename, fps=10):
|
|
|
b, c, t, h, w = x.shape
|
|
|
|
|
|
per_row = b
|
|
|
for p in [6, 5, 4, 3, 2]:
|
|
|
if b % p == 0:
|
|
|
per_row = p
|
|
|
break
|
|
|
|
|
|
os.makedirs(os.path.dirname(os.path.abspath(os.path.realpath(output_filename))), exist_ok=True)
|
|
|
x = torch.clamp(x.float(), -1., 1.) * 127.5 + 127.5
|
|
|
x = x.detach().cpu().to(torch.uint8)
|
|
|
x = einops.rearrange(x, '(m n) c t h w -> t (m h) (n w) c', n=per_row)
|
|
|
torchvision.io.write_video(output_filename, x, fps=fps, video_codec='libx264', options={'crf': '0'})
|
|
|
return x
|
|
|
|
|
|
|
|
|
def save_bcthw_as_png(x, output_filename):
|
|
|
os.makedirs(os.path.dirname(os.path.abspath(os.path.realpath(output_filename))), exist_ok=True)
|
|
|
x = torch.clamp(x.float(), -1., 1.) * 127.5 + 127.5
|
|
|
x = x.detach().cpu().to(torch.uint8)
|
|
|
x = einops.rearrange(x, 'b c t h w -> c (b h) (t w)')
|
|
|
torchvision.io.write_png(x, output_filename)
|
|
|
return output_filename
|
|
|
|
|
|
|
|
|
def save_bchw_as_png(x, output_filename):
|
|
|
os.makedirs(os.path.dirname(os.path.abspath(os.path.realpath(output_filename))), exist_ok=True)
|
|
|
x = torch.clamp(x.float(), -1., 1.) * 127.5 + 127.5
|
|
|
x = x.detach().cpu().to(torch.uint8)
|
|
|
x = einops.rearrange(x, 'b c h w -> c h (b w)')
|
|
|
torchvision.io.write_png(x, output_filename)
|
|
|
return output_filename
|
|
|
|
|
|
|
|
|
def add_tensors_with_padding(tensor1, tensor2):
|
|
|
if tensor1.shape == tensor2.shape:
|
|
|
return tensor1 + tensor2
|
|
|
|
|
|
shape1 = tensor1.shape
|
|
|
shape2 = tensor2.shape
|
|
|
|
|
|
new_shape = tuple(max(s1, s2) for s1, s2 in zip(shape1, shape2))
|
|
|
|
|
|
padded_tensor1 = torch.zeros(new_shape)
|
|
|
padded_tensor2 = torch.zeros(new_shape)
|
|
|
|
|
|
padded_tensor1[tuple(slice(0, s) for s in shape1)] = tensor1
|
|
|
padded_tensor2[tuple(slice(0, s) for s in shape2)] = tensor2
|
|
|
|
|
|
result = padded_tensor1 + padded_tensor2
|
|
|
return result
|
|
|
|
|
|
|
|
|
def print_free_mem():
|
|
|
torch.cuda.empty_cache()
|
|
|
free_mem, total_mem = torch.cuda.mem_get_info(0)
|
|
|
free_mem_mb = free_mem / (1024 ** 2)
|
|
|
total_mem_mb = total_mem / (1024 ** 2)
|
|
|
print(f"Free memory: {free_mem_mb:.2f} MB")
|
|
|
print(f"Total memory: {total_mem_mb:.2f} MB")
|
|
|
return
|
|
|
|
|
|
|
|
|
def print_gpu_parameters(device, state_dict, log_count=1):
|
|
|
summary = {"device": device, "keys_count": len(state_dict)}
|
|
|
|
|
|
logged_params = {}
|
|
|
for i, (key, tensor) in enumerate(state_dict.items()):
|
|
|
if i >= log_count:
|
|
|
break
|
|
|
logged_params[key] = tensor.flatten()[:3].tolist()
|
|
|
|
|
|
summary["params"] = logged_params
|
|
|
|
|
|
print(str(summary))
|
|
|
return
|
|
|
|
|
|
|
|
|
def visualize_txt_as_img(width, height, text, font_path='font/DejaVuSans.ttf', size=18):
|
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
|
|
|
|
txt = Image.new("RGB", (width, height), color="white")
|
|
|
draw = ImageDraw.Draw(txt)
|
|
|
font = ImageFont.truetype(font_path, size=size)
|
|
|
|
|
|
if text == '':
|
|
|
return np.array(txt)
|
|
|
|
|
|
|
|
|
lines = []
|
|
|
words = text.split()
|
|
|
current_line = words[0]
|
|
|
|
|
|
for word in words[1:]:
|
|
|
line_with_word = f"{current_line} {word}"
|
|
|
if draw.textbbox((0, 0), line_with_word, font=font)[2] <= width:
|
|
|
current_line = line_with_word
|
|
|
else:
|
|
|
lines.append(current_line)
|
|
|
current_line = word
|
|
|
|
|
|
lines.append(current_line)
|
|
|
|
|
|
|
|
|
y = 0
|
|
|
line_height = draw.textbbox((0, 0), "A", font=font)[3]
|
|
|
|
|
|
for line in lines:
|
|
|
if y + line_height > height:
|
|
|
break
|
|
|
draw.text((0, y), line, fill="black", font=font)
|
|
|
y += line_height
|
|
|
|
|
|
return np.array(txt)
|
|
|
|
|
|
|
|
|
def blue_mark(x):
|
|
|
x = x.copy()
|
|
|
c = x[:, :, 2]
|
|
|
b = cv2.blur(c, (9, 9))
|
|
|
x[:, :, 2] = ((c - b) * 16.0 + b).clip(-1, 1)
|
|
|
return x
|
|
|
|
|
|
|
|
|
def green_mark(x):
|
|
|
x = x.copy()
|
|
|
x[:, :, 2] = -1
|
|
|
x[:, :, 0] = -1
|
|
|
return x
|
|
|
|
|
|
|
|
|
def frame_mark(x):
|
|
|
x = x.copy()
|
|
|
x[:64] = -1
|
|
|
x[-64:] = -1
|
|
|
x[:, :8] = 1
|
|
|
x[:, -8:] = 1
|
|
|
return x
|
|
|
|
|
|
|
|
|
@torch.inference_mode()
|
|
|
def pytorch2numpy(imgs):
|
|
|
results = []
|
|
|
for x in imgs:
|
|
|
y = x.movedim(0, -1)
|
|
|
y = y * 127.5 + 127.5
|
|
|
y = y.detach().float().cpu().numpy().clip(0, 255).astype(np.uint8)
|
|
|
results.append(y)
|
|
|
return results
|
|
|
|
|
|
|
|
|
@torch.inference_mode()
|
|
|
def numpy2pytorch(imgs):
|
|
|
h = torch.from_numpy(np.stack(imgs, axis=0)).float() / 127.5 - 1.0
|
|
|
h = h.movedim(-1, 1)
|
|
|
return h
|
|
|
|
|
|
|
|
|
@torch.no_grad()
|
|
|
def duplicate_prefix_to_suffix(x, count, zero_out=False):
|
|
|
if zero_out:
|
|
|
return torch.cat([x, torch.zeros_like(x[:count])], dim=0)
|
|
|
else:
|
|
|
return torch.cat([x, x[:count]], dim=0)
|
|
|
|
|
|
|
|
|
def weighted_mse(a, b, weight):
|
|
|
return torch.mean(weight.float() * (a.float() - b.float()) ** 2)
|
|
|
|
|
|
|
|
|
def clamped_linear_interpolation(x, x_min, y_min, x_max, y_max, sigma=1.0):
|
|
|
x = (x - x_min) / (x_max - x_min)
|
|
|
x = max(0.0, min(x, 1.0))
|
|
|
x = x ** sigma
|
|
|
return y_min + x * (y_max - y_min)
|
|
|
|
|
|
|
|
|
def expand_to_dims(x, target_dims):
|
|
|
return x.view(*x.shape, *([1] * max(0, target_dims - x.dim())))
|
|
|
|
|
|
|
|
|
def repeat_to_batch_size(tensor: torch.Tensor, batch_size: int):
|
|
|
if tensor is None:
|
|
|
return None
|
|
|
|
|
|
first_dim = tensor.shape[0]
|
|
|
|
|
|
if first_dim == batch_size:
|
|
|
return tensor
|
|
|
|
|
|
if batch_size % first_dim != 0:
|
|
|
raise ValueError(f"Cannot evenly repeat first dim {first_dim} to match batch_size {batch_size}.")
|
|
|
|
|
|
repeat_times = batch_size // first_dim
|
|
|
|
|
|
return tensor.repeat(repeat_times, *[1] * (tensor.dim() - 1))
|
|
|
|
|
|
|
|
|
def dim5(x):
|
|
|
return expand_to_dims(x, 5)
|
|
|
|
|
|
|
|
|
def dim4(x):
|
|
|
return expand_to_dims(x, 4)
|
|
|
|
|
|
|
|
|
def dim3(x):
|
|
|
return expand_to_dims(x, 3)
|
|
|
|
|
|
|
|
|
def crop_or_pad_yield_mask(x, length):
|
|
|
B, F, C = x.shape
|
|
|
device = x.device
|
|
|
dtype = x.dtype
|
|
|
|
|
|
if F < length:
|
|
|
y = torch.zeros((B, length, C), dtype=dtype, device=device)
|
|
|
mask = torch.zeros((B, length), dtype=torch.bool, device=device)
|
|
|
y[:, :F, :] = x
|
|
|
mask[:, :F] = True
|
|
|
return y, mask
|
|
|
|
|
|
return x[:, :length, :], torch.ones((B, length), dtype=torch.bool, device=device)
|
|
|
|
|
|
|
|
|
def extend_dim(x, dim, minimal_length, zero_pad=False):
|
|
|
original_length = int(x.shape[dim])
|
|
|
|
|
|
if original_length >= minimal_length:
|
|
|
return x
|
|
|
|
|
|
if zero_pad:
|
|
|
padding_shape = list(x.shape)
|
|
|
padding_shape[dim] = minimal_length - original_length
|
|
|
padding = torch.zeros(padding_shape, dtype=x.dtype, device=x.device)
|
|
|
else:
|
|
|
idx = (slice(None),) * dim + (slice(-1, None),) + (slice(None),) * (len(x.shape) - dim - 1)
|
|
|
last_element = x[idx]
|
|
|
padding = last_element.repeat_interleave(minimal_length - original_length, dim=dim)
|
|
|
|
|
|
return torch.cat([x, padding], dim=dim)
|
|
|
|
|
|
|
|
|
def lazy_positional_encoding(t, repeats=None):
|
|
|
if not isinstance(t, list):
|
|
|
t = [t]
|
|
|
|
|
|
from diffusers.models.embeddings import get_timestep_embedding
|
|
|
|
|
|
te = torch.tensor(t)
|
|
|
te = get_timestep_embedding(timesteps=te, embedding_dim=256, flip_sin_to_cos=True, downscale_freq_shift=0.0, scale=1.0)
|
|
|
|
|
|
if repeats is None:
|
|
|
return te
|
|
|
|
|
|
te = te[:, None, :].expand(-1, repeats, -1)
|
|
|
|
|
|
return te
|
|
|
|
|
|
|
|
|
def state_dict_offset_merge(A, B, C=None):
|
|
|
result = {}
|
|
|
keys = A.keys()
|
|
|
|
|
|
for key in keys:
|
|
|
A_value = A[key]
|
|
|
B_value = B[key].to(A_value)
|
|
|
|
|
|
if C is None:
|
|
|
result[key] = A_value + B_value
|
|
|
else:
|
|
|
C_value = C[key].to(A_value)
|
|
|
result[key] = A_value + B_value - C_value
|
|
|
|
|
|
return result
|
|
|
|
|
|
|
|
|
def state_dict_weighted_merge(state_dicts, weights):
|
|
|
if len(state_dicts) != len(weights):
|
|
|
raise ValueError("Number of state dictionaries must match number of weights")
|
|
|
|
|
|
if not state_dicts:
|
|
|
return {}
|
|
|
|
|
|
total_weight = sum(weights)
|
|
|
|
|
|
if total_weight == 0:
|
|
|
raise ValueError("Sum of weights cannot be zero")
|
|
|
|
|
|
normalized_weights = [w / total_weight for w in weights]
|
|
|
|
|
|
keys = state_dicts[0].keys()
|
|
|
result = {}
|
|
|
|
|
|
for key in keys:
|
|
|
result[key] = state_dicts[0][key] * normalized_weights[0]
|
|
|
|
|
|
for i in range(1, len(state_dicts)):
|
|
|
state_dict_value = state_dicts[i][key].to(result[key])
|
|
|
result[key] += state_dict_value * normalized_weights[i]
|
|
|
|
|
|
return result
|
|
|
|
|
|
|
|
|
def group_files_by_folder(all_files):
|
|
|
grouped_files = {}
|
|
|
|
|
|
for file in all_files:
|
|
|
folder_name = os.path.basename(os.path.dirname(file))
|
|
|
if folder_name not in grouped_files:
|
|
|
grouped_files[folder_name] = []
|
|
|
grouped_files[folder_name].append(file)
|
|
|
|
|
|
list_of_lists = list(grouped_files.values())
|
|
|
return list_of_lists
|
|
|
|
|
|
|
|
|
def generate_timestamp():
|
|
|
now = datetime.datetime.now()
|
|
|
timestamp = now.strftime('%y%m%d_%H%M%S')
|
|
|
milliseconds = f"{int(now.microsecond / 1000):03d}"
|
|
|
random_number = random.randint(0, 9999)
|
|
|
return f"{timestamp}_{milliseconds}_{random_number}"
|
|
|
|
|
|
|
|
|
def write_PIL_image_with_png_info(image, metadata, path):
|
|
|
from PIL.PngImagePlugin import PngInfo
|
|
|
|
|
|
png_info = PngInfo()
|
|
|
for key, value in metadata.items():
|
|
|
png_info.add_text(key, value)
|
|
|
|
|
|
image.save(path, "PNG", pnginfo=png_info)
|
|
|
return image
|
|
|
|
|
|
|
|
|
def torch_safe_save(content, path):
|
|
|
torch.save(content, path + '_tmp')
|
|
|
os.replace(path + '_tmp', path)
|
|
|
return path
|
|
|
|
|
|
|
|
|
def move_optimizer_to_device(optimizer, device):
|
|
|
for state in optimizer.state.values():
|
|
|
for k, v in state.items():
|
|
|
if isinstance(v, torch.Tensor):
|
|
|
state[k] = v.to(device)
|
|
|
|
