Spaces:
Sleeping
Sleeping
File size: 6,254 Bytes
9d53ea8 2fe2169 9d53ea8 2fe2169 9d53ea8 2fe2169 9d53ea8 | 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 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 | # Based on algorithm 5 but in latent space of TAESD
import torch
import torch.nn.functional as F
from torch import nn
from tqdm import tqdm
import numpy as np
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
from algorithms.taesd.taesd import TAESD
PATCHSIZE = 12
STRIDE = 6
ITERATIONS = 50
INITIAL_NOISE_FACTOR = 0.5
INIT_PC = 1
INIT_PC_FACTOR = 0.5
ALGO_SAMPLES = 500
SEED = None
OUTPUT_RGB_GIF = True
OUTPUT_LATENT_GIF = False
@torch.no_grad()
def extract_centered_patches(img, patchsize, stride:int = 1):
"""
Extrait les patches centrés pour chaque pixel de l'image.
"""
img_padded = img
return F.unfold(img_padded, kernel_size=patchsize, padding=0, stride=stride)
@torch.no_grad()
def weighted_patch_average(P_synth, patchsize,C, H, W, mode="standard", device='cpu', stride: int = 1):
if mode == "gaussian":
# Gaussian weight
half_win_size = patchsize // 2
sig_pix = 1.0 * half_win_size
# arange creates the spatial spread
dx = torch.linspace(-half_win_size, half_win_size, steps=patchsize, device=device)
w1d = torch.exp(-(dx / sig_pix) ** 2 / 2.0)
# Create 2D weight and adapt to 3 color channels
w2d = w1d[:, None] * w1d[None, :]
w = w2d.repeat(C, 1, 1).reshape(-1)
w = w / w.sum() # normalize
w = w.unsqueeze(0).unsqueeze(-1) # (1, C * patchsize^2, 1)
else: # standard
# NIFTY weight
w=torch.exp(-torch.linspace(-patchsize//2,patchsize//2,steps=patchsize).pow(2)/2/(patchsize*1/4)**2).to(device)
w = w.view(-1, 1) * w.view(1, -1)
w = w.repeat(C, 1, 1).view(-1) #
w /= w.sum() # Normalization
w = w.unsqueeze(0).unsqueeze(-1)
fold_layer = nn.Fold((W, H), kernel_size=patchsize, dilation=1, padding=0, stride=stride)
# Apply weights and fold
synth = fold_layer(P_synth * w)
count = fold_layer(P_synth * 0 + w)
count= (count*(count!=0)+1.*(count==0))
synth = synth / count
return synth
def make_times(n_timestep, schedule='cosine', t0=0):
'''
different time discretizations (0 to 1), 'quad' has smaller timesteps near t=0
'''
times = torch.linspace(t0, 1., n_timestep + 1) # default fallback to linear
if schedule == "linear":
times = torch.linspace(t0, 1., n_timestep + 1)
elif schedule == "quad":
times = torch.linspace(t0 ** 0.5, 1., n_timestep + 1) ** 2
elif schedule == "cosine":
times = (
torch.linspace(torch.arcsin(torch.tensor(t0) ** 0.5), torch.pi / 2, n_timestep + 1)
)
times = torch.sin(times).pow(2)
times = times / times[-1]
return times
@torch.no_grad()
def ALGO_LATENT(clean_batch_tensor, initialisation_tensor, patchsize=3, N=50, schedule='linear', device='cpu', mask_weight_type="standard", seed=None, stride:int=1):
if seed is not None:
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
clean_batch_tensor = clean_batch_tensor.to(device)
N_imgs, C, H, W = clean_batch_tensor.shape
# Extract patches from reference dataset
Z = extract_centered_patches(clean_batch_tensor, patchsize, stride=stride)
x_noise = initialisation_tensor
x_noise = (1. - INITIAL_NOISE_FACTOR ) * x_noise + INITIAL_NOISE_FACTOR * torch.randn_like(x_noise, device=device)
#t0 = (1. - INITIAL_NOISE_FACTOR )
t0=0.
x_n1 = x_noise.clone()
saved_steps = [x_n1.cpu()]
times = make_times(N, schedule, t0)
for it in tqdm(range(N)):
t = times[it]
delta_t = times[it + 1] - t
x_patches = extract_centered_patches(x_n1, patchsize, stride=stride) # (1, C*patchsize^2, H*W)
dists = torch.sum((x_patches - Z * t)**2, dim=1) # (N_imgs, H*W)
w = torch.softmax(-dists / (2 * ((1. - t) ** 2)), dim=0)# (N_imgs, H*W)
v = ((Z - x_patches) * w.unsqueeze(1)).sum(0, keepdim=True) / (1. - t)
# Z : (N_imgs, C*patchsize^2, H*W) // x_patches (1, C*patchsize^2, H*W) // w (N_imgs, H*W)
x_patches_updated = x_patches + v * delta_t
x_n1 = weighted_patch_average(x_patches_updated, patchsize, C, H, W, mode=mask_weight_type, device=device, stride=stride)
yield x_n1.cpu()
saved_steps.append(x_n1.cpu())
return saved_steps
@torch.no_grad()
def imgs_to_gif(path="out.gif",imgs=None):
from PIL import Image
np_imgs = [np.uint8(np.clip(img.permute(0, 2, 3, 1).detach().numpy()[0] , 0, 1) * 255) for img in imgs]
im_list = [Image.fromarray(np_img, mode='RGB') for np_img in np_imgs]
im_list = im_list + [im_list[-1]] * 10
im_list[0].save(path, save_all=True, append_images=im_list[1:], duration=5, loop=0)
@torch.no_grad()
def imgs_to_gif_encode(imgs=None, device='cuda:0'):
from PIL import Image
imgs_encode_to_rgb = [ TAESD.scale_latents(i.to(device))[0][:3].permute(1,2,0).detach().cpu().numpy() for i in imgs]
np_imgs = [np.uint8(img * 255) for img in imgs_encode_to_rgb]
im_list = [Image.fromarray(np_img, mode='RGB') for np_img in np_imgs]
im_list = im_list + [im_list[-1]] * 100
im_list[0].save("out_a5_enc.gif", save_all=True, append_images=im_list[1:], duration=1, loop=0)
@torch.no_grad()
def transform_to_latent_space(tensors, taesd:TAESD, dev):
from PIL import Image
import torchvision.transforms.functional as TF
# Scale latents immediately after encoding
encoded_tensor = taesd.encoder(tensors[0].unsqueeze(0).to(dev))
cat_tensor = torch.zeros(tensors.shape[0], *encoded_tensor.shape[1:]).to(dev)
cat_tensor[0] = encoded_tensor
for i in tqdm(range(1, tensors.shape[0]), "Encoding tensors"):
encoded_tensor = taesd.encoder(tensors[i].unsqueeze(0).to(dev))
cat_tensor[i] = encoded_tensor
return cat_tensor
@torch.no_grad()
def transform_to_latent_space_tensor(tensor, taesd:TAESD, dev):
import torchvision.transforms.functional as TF
encoded_tensor = taesd.encoder(tensor.to(dev))
return encoded_tensor
@torch.no_grad()
def decode_tensor(tensor:torch.TensorType, taesd:TAESD):
return taesd.decoder(tensor).clamp(0, 1)
|