73 lines
2.2 KiB
Python
73 lines
2.2 KiB
Python
import math
|
|
import numpy as np
|
|
from PIL import Image
|
|
import torch
|
|
import torch.nn as nn
|
|
import importlib
|
|
from .model.utils import get_named_beta_schedule, _extract_into_tensor
|
|
from copy import deepcopy
|
|
|
|
|
|
|
|
def prepare_mask(mask):
|
|
mask = mask.float()[0]
|
|
old_mask = deepcopy(mask)
|
|
for i in range(mask.shape[1]):
|
|
for j in range(mask.shape[2]):
|
|
if old_mask[0][i][j] == 1:
|
|
continue
|
|
if i != 0:
|
|
mask[:, i - 1, j] = 0
|
|
if j != 0:
|
|
mask[:, i, j - 1] = 0
|
|
if i != 0 and j != 0:
|
|
mask[:, i - 1, j - 1] = 0
|
|
if i != mask.shape[1] - 1:
|
|
mask[:, i + 1, j] = 0
|
|
if j != mask.shape[2] - 1:
|
|
mask[:, i, j + 1] = 0
|
|
if i != mask.shape[1] - 1 and j != mask.shape[2] - 1:
|
|
mask[:, i + 1, j + 1] = 0
|
|
return mask.unsqueeze(0)
|
|
|
|
|
|
def prepare_image(pil_image, w=512, h=512):
|
|
pil_image = pil_image.resize((w, h), resample=Image.BICUBIC, reducing_gap=1)
|
|
arr = np.array(pil_image.convert("RGB"))
|
|
arr = arr.astype(np.float32) / 127.5 - 1
|
|
arr = np.transpose(arr, [2, 0, 1])
|
|
image = torch.from_numpy(arr).unsqueeze(0)
|
|
return image
|
|
|
|
|
|
def q_sample(x_start, t, schedule_name="linear", num_steps=1000, noise=None):
|
|
betas = get_named_beta_schedule(schedule_name, num_steps)
|
|
alphas = 1.0 - betas
|
|
alphas_cumprod = np.cumprod(alphas, axis=0)
|
|
sqrt_alphas_cumprod = np.sqrt(alphas_cumprod)
|
|
sqrt_one_minus_alphas_cumprod = np.sqrt(1.0 - alphas_cumprod)
|
|
if noise is None:
|
|
noise = torch.randn_like(x_start)
|
|
assert noise.shape == x_start.shape
|
|
return (
|
|
_extract_into_tensor(sqrt_alphas_cumprod, t, x_start.shape) * x_start
|
|
+ _extract_into_tensor(sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
|
|
)
|
|
|
|
|
|
def process_images(batch):
|
|
scaled = (
|
|
((batch + 1) * 127.5)
|
|
.round()
|
|
.clamp(0, 255)
|
|
.to(torch.uint8)
|
|
.to("cpu")
|
|
.permute(0, 2, 3, 1)
|
|
.numpy()
|
|
)
|
|
images = []
|
|
for i in range(scaled.shape[0]):
|
|
images.append(Image.fromarray(scaled[i]))
|
|
return images
|
|
|