| 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 | |
| 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( torchvision.transforms.functional.to_pil_image(img.clamp(0, 1))) | |
| 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 ) | |