231 lines
6.3 KiB
Python
231 lines
6.3 KiB
Python
import PIL
|
|
import torch
|
|
import requests
|
|
import torchvision
|
|
from math import ceil
|
|
from io import BytesIO
|
|
import matplotlib.pyplot as plt
|
|
import torchvision.transforms.functional as F
|
|
import math
|
|
from tqdm import tqdm
|
|
import numpy as np
|
|
|
|
|
|
def download_image(url):
|
|
return PIL.Image.open(requests.get(url, stream=True).raw).convert("RGB")
|
|
|
|
|
|
def resize_image(image, size=768):
|
|
tensor_image = F.to_tensor(image)
|
|
resized_image = F.resize(tensor_image, size, antialias=True)
|
|
return resized_image
|
|
|
|
|
|
def downscale_images(images, factor=3 / 4):
|
|
scaled_height, scaled_width = int(((images.size(-2) * factor) // 32) * 32), int(
|
|
((images.size(-1) * factor) // 32) * 32
|
|
)
|
|
scaled_image = torchvision.transforms.functional.resize(
|
|
images,
|
|
(scaled_height, scaled_width),
|
|
interpolation=torchvision.transforms.InterpolationMode.NEAREST,
|
|
)
|
|
return scaled_image
|
|
|
|
|
|
def calculate_latent_sizes(
|
|
height=1024,
|
|
width=1024,
|
|
batch_size=4,
|
|
compression_factor_b=42.67,
|
|
compression_factor_a=4.0,
|
|
):
|
|
resolution_multiple = 42.67
|
|
latent_height = ceil(height / compression_factor_b)
|
|
latent_width = ceil(width / compression_factor_b)
|
|
stage_c_latent_shape = (batch_size, 16, latent_height, latent_width)
|
|
|
|
latent_height = ceil(height / compression_factor_a)
|
|
latent_width = ceil(width / compression_factor_a)
|
|
stage_b_latent_shape = (batch_size, 4, latent_height, latent_width)
|
|
|
|
return stage_c_latent_shape, stage_b_latent_shape
|
|
|
|
|
|
def get_views(H, W, window_size=64, stride=16):
|
|
"""
|
|
- H, W: height and width of the latent
|
|
"""
|
|
num_blocks_height = (H - window_size) // stride + 1
|
|
num_blocks_width = (W - window_size) // stride + 1
|
|
total_num_blocks = int(num_blocks_height * num_blocks_width)
|
|
views = []
|
|
for i in range(total_num_blocks):
|
|
h_start = int((i // num_blocks_width) * stride)
|
|
h_end = h_start + window_size
|
|
w_start = int((i % num_blocks_width) * stride)
|
|
w_end = w_start + window_size
|
|
views.append((h_start, h_end, w_start, w_end))
|
|
return views
|
|
|
|
|
|
def show_images(images, rows=None, cols=None, **kwargs):
|
|
if images.size(1) == 1:
|
|
images = images.repeat(1, 3, 1, 1)
|
|
elif images.size(1) > 3:
|
|
images = images[:, :3]
|
|
|
|
if rows is None:
|
|
rows = 1
|
|
if cols is None:
|
|
cols = images.size(0) // rows
|
|
|
|
_, _, h, w = images.shape
|
|
|
|
imgs = []
|
|
for i, img in enumerate(images):
|
|
imgs.append(
|
|
torch.from_numpy(
|
|
np.array(
|
|
torchvision.transforms.functional.to_pil_image(
|
|
img.clamp(0, 1)
|
|
).convert("RGB")
|
|
).astype(np.float32)
|
|
/ 255.0
|
|
)[
|
|
None,
|
|
]
|
|
)
|
|
|
|
return imgs
|
|
|
|
|
|
def decode_b(
|
|
conditions_b,
|
|
unconditions_b,
|
|
models_b,
|
|
bshape,
|
|
extras_b,
|
|
device,
|
|
stage_a_tiled=False,
|
|
num_instance=4,
|
|
patch_size=256,
|
|
stride=24,
|
|
):
|
|
|
|
sampling_b = extras_b.gdf.sample(
|
|
models_b.generator.half(),
|
|
conditions_b,
|
|
bshape,
|
|
unconditions_b,
|
|
device=device,
|
|
**extras_b.sampling_configs,
|
|
)
|
|
models_b.generator.cuda()
|
|
for sampled_b, _, _ in tqdm(
|
|
sampling_b, total=extras_b.sampling_configs["timesteps"]
|
|
):
|
|
sampled_b = sampled_b
|
|
models_b.generator.cpu()
|
|
torch.cuda.empty_cache()
|
|
if stage_a_tiled:
|
|
with torch.cuda.amp.autocast(dtype=torch.float16):
|
|
padding = (stride * 2, stride * 2, stride * 2, stride * 2)
|
|
sampled_b = torch.nn.functional.pad(sampled_b, padding, mode="reflect")
|
|
count = torch.zeros(
|
|
(
|
|
sampled_b.shape[0],
|
|
3,
|
|
sampled_b.shape[-2] * 4,
|
|
sampled_b.shape[-1] * 4,
|
|
),
|
|
requires_grad=False,
|
|
device=sampled_b.device,
|
|
)
|
|
sampled = torch.zeros(
|
|
(
|
|
sampled_b.shape[0],
|
|
3,
|
|
sampled_b.shape[-2] * 4,
|
|
sampled_b.shape[-1] * 4,
|
|
),
|
|
requires_grad=False,
|
|
device=sampled_b.device,
|
|
)
|
|
views = get_views(
|
|
sampled_b.shape[-2],
|
|
sampled_b.shape[-1],
|
|
window_size=patch_size,
|
|
stride=stride,
|
|
)
|
|
|
|
for view_idx, (h_start, h_end, w_start, w_end) in enumerate(
|
|
tqdm(views, total=len(views))
|
|
):
|
|
|
|
sampled[
|
|
:, :, h_start * 4 : h_end * 4, w_start * 4 : w_end * 4
|
|
] += models_b.stage_a.decode(
|
|
sampled_b[:, :, h_start:h_end, w_start:w_end]
|
|
).float()
|
|
count[:, :, h_start * 4 : h_end * 4, w_start * 4 : w_end * 4] += 1
|
|
sampled /= count
|
|
sampled = sampled[
|
|
:, :, stride * 4 * 2 : -stride * 4 * 2, stride * 4 * 2 : -stride * 4 * 2
|
|
]
|
|
else:
|
|
|
|
sampled = models_b.stage_a.decode(sampled_b, tiled_decoding=stage_a_tiled)
|
|
|
|
return sampled.float()
|
|
|
|
|
|
def generation_c(
|
|
batch,
|
|
models,
|
|
extras,
|
|
core,
|
|
stage_c_latent_shape,
|
|
stage_c_latent_shape_lr,
|
|
device,
|
|
conditions=None,
|
|
unconditions=None,
|
|
):
|
|
if conditions is None:
|
|
conditions = core.get_conditions(
|
|
batch,
|
|
models,
|
|
extras,
|
|
is_eval=True,
|
|
is_unconditional=False,
|
|
eval_image_embeds=False,
|
|
)
|
|
if unconditions is None:
|
|
unconditions = core.get_conditions(
|
|
batch,
|
|
models,
|
|
extras,
|
|
is_eval=True,
|
|
is_unconditional=True,
|
|
eval_image_embeds=False,
|
|
)
|
|
sampling_c = extras.gdf.sample(
|
|
models.generator,
|
|
conditions,
|
|
stage_c_latent_shape,
|
|
stage_c_latent_shape_lr,
|
|
unconditions,
|
|
device=device,
|
|
**extras.sampling_configs,
|
|
)
|
|
for idx, (sampled_c, sampled_c_curr, _, _) in enumerate(
|
|
tqdm(sampling_c, total=extras.sampling_configs["timesteps"])
|
|
):
|
|
sampled_c = sampled_c
|
|
return sampled_c
|
|
|
|
|
|
def get_target_lr_size(ratio, std_size=24):
|
|
w, h = int(std_size / math.sqrt(ratio)), int(std_size * math.sqrt(ratio))
|
|
return (h * 32, w * 32)
|