File size: 5,009 Bytes
746c807 |
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 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 |
"""Utility functions"""
import importlib
import random
import torch
import numpy as np
from PIL import Image
class UnNormalize(object):
"""Unformalize image as: image = (image * std) + mean
"""
def __init__(self, mean, std):
self.mean = torch.tensor(mean)
self.std = torch.tensor(std)
def __call__(self, tensor):
"""
Args:
tensor: A tensor of shape [C, H, W] or [N, C, H, W]
Returns:
tensor: A tensor of shape [C, H, W] or [N, C, H, W]
"""
std = self.std.to(tensor.device)
mean = self.mean.to(tensor.device)
if tensor.ndim == 3:
std, mean = std.view(-1, 1, 1), mean.view(-1, 1, 1)
elif tensor.ndim == 4:
std, mean = std.view(1, -1, 1, 1), mean.view(1, -1, 1, 1)
tensor = (tensor * std) + mean
return tensor
class VQVAEUnNormalize(UnNormalize):
"""Unformalize image as:
First: image = (image * std) + mean
Second: image = (image * 2) - 1
"""
def __call__(self, tensor):
"""
Args:
tensor (Tensor): Tensor image of size (C, H, W) or (N, C, H, W)
to be unnormalized.
Returns:
Tensor: UnNormalized image.
"""
tensor = super().__call__(tensor)
tensor = 2 * tensor - 1
return tensor
def normalize(image,rescale=True):
if rescale:
image = image.float() / 255.0 # Convert to float and rescale to [0, 1]
normalize_image = 2*image-1 # normalize to [-1, 1]
return normalize_image
# train_transforms = transforms.Compose(
# [
# transforms.Resize(args.resolution, interpolation=transforms.InterpolationMode.BILINEAR),
# transforms.CenterCrop(args.resolution) if args.center_crop else transforms.RandomCrop(args.resolution),
# transforms.RandomHorizontalFlip() if args.random_flip else transforms.Lambda(lambda x: x),
# transforms.ToTensor(),
# transforms.Normalize([0.5], [0.5]),
# ]
# )
def mean_list(l):
l = [int(_l) for _l in l]
return float(sum(l)) / len(l)
def segment_mean(x, index):
"""Function as tf.segment_mean.
"""
x = x.view(-1, x.shape[-1])
index = index.view(-1)
max_index = index.max() + 1
sum_x = torch.zeros((max_index, x.shape[-1]),
dtype=x.dtype,
device=x.device)
num_index = torch.zeros((max_index,),
dtype=x.dtype,
device=x.device)
num_index = num_index.scatter_add_(
0, index, torch.ones_like(index, dtype=x.dtype))
num_index = torch.where(torch.eq(num_index, 0),
torch.ones_like(num_index, dtype=x.dtype),
num_index)
index_2d = index.view(-1, 1).expand(-1, x.shape[-1])
sum_x = sum_x.scatter_add_(0, index_2d, x)
mean_x = sum_x.div_(num_index.view(-1, 1))
return mean_x
def initiate_time_steps(step, total_timestep, batch_size, config):
"""A helper function to initiate time steps for the diffusion model.
Args:
step: An integer of the constant step
total_timestep: An integer of the total timesteps of the diffusion model
batch_size: An integer of the batch size
config: A config object
Returns:
timesteps: A tensor of shape [batch_size,] of the time steps
"""
if config.tta.rand_timestep_equal_int:
# the same timestep for each image in the batch
interval_val = total_timestep // batch_size
start_point = random.randint(0, interval_val - 1)
timesteps = torch.tensor(
list(range(start_point, total_timestep, interval_val))
).long()
return timesteps
elif config.tta.random_timestep_per_iteration:
# random timestep for each image in the batch
return torch.randint(0, total_timestep, (batch_size,)).long() #default
else:
# why we need to do this?
return torch.tensor([step] * batch_size).long()
def instantiate_from_config(config):
"""A helper function to instantiate a class from a config object.
See https://github.com/CompVis/stable-diffusion/blob/main/ldm/util.py
"""
if not "target" in config:
if config == '__is_first_stage__':
return None
elif config == "__is_unconditional__":
return None
raise KeyError("Expected key `target` to instantiate.")
return get_obj_from_str(config["target"])(**config.get("params", dict()))
def get_obj_from_str(string, reload=False):
"""A helper function to instantiate a class from a config object.
See https://github.com/CompVis/stable-diffusion/blob/main/ldm/util.py
"""
module, cls = string.rsplit(".", 1)
if reload:
module_imp = importlib.import_module(module)
importlib.reload(module_imp)
return getattr(importlib.import_module(module, package=None), cls) |