import torch
from ldm.models.diffusion.ddim import DDIMSampler
from ldm.models.diffusion.plms import PLMSSampler
from ldm.util import instantiate_from_config
import numpy as np
import random
from dataset.concat_dataset import ConCatDataset #, collate_fn
from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
import os
from tqdm import tqdm
from distributed import get_rank, synchronize, get_world_size
from trainer import read_official_ckpt, batch_to_device, ImageCaptionSaver, wrap_loader #, get_padded_boxes
from PIL import Image
import math
import json
def draw_masks_from_boxes(boxes,size):
image_masks = []
for box in boxes:
image_mask = torch.ones(size[0],size[1])
for bx in box:
x0, x1 = bx[0]*size[0], bx[2]*size[0]
y0, y1 = bx[1]*size[1], bx[3]*size[1]
image_mask[int(y0):int(y1), int(x0):int(x1)] = 0
image_masks.append(image_mask)
return torch.stack(image_masks).unsqueeze(1)
def set_alpha_scale(model, alpha_scale):
from ldm.modules.attention import GatedCrossAttentionDense, GatedSelfAttentionDense
for module in model.modules():
if type(module) == GatedCrossAttentionDense or type(module) == GatedSelfAttentionDense:
module.scale = alpha_scale
# print("scale: ", alpha_scale)
# print("attn: ", module.alpha_attn)
# print("dense: ", module.alpha_dense)
# print(' ')
# print(' ')
def save_images(samples, image_ids, folder, to256):
for sample, image_id in zip(samples, image_ids):
sample = torch.clamp(sample, min=-1, max=1) * 0.5 + 0.5
sample = sample.cpu().numpy().transpose(1,2,0) * 255
img_name = str(int(image_id))+'.png'
img = Image.fromarray(sample.astype(np.uint8))
if to256:
img = img.resize( (256,256), Image.BICUBIC)
img.save(os.path.join(folder,img_name))
def ckpt_to_folder_name(basename):
name=""
for s in basename:
if s.isdigit():
name+=s
seen = round( int(name)/1000, 1 )
return str(seen).ljust(4,'0')+'k'
class Evaluator:
def __init__(self, config):
self.config = config
self.device = torch.device("cuda")
# = = = = = create model and diffusion = = = = = #
if self.config.ckpt != "real":
self.model = instantiate_from_config(config.model).to(self.device)
self.autoencoder = instantiate_from_config(config.autoencoder).to(self.device)
self.text_encoder = instantiate_from_config(config.text_encoder).to(self.device)
self.diffusion = instantiate_from_config(config.diffusion).to(self.device)
# donot need to load official_ckpt for self.model here, since we will load from our ckpt
state_dict = read_official_ckpt( os.path.join(config.DATA_ROOT, config.official_ckpt_name) )
self.autoencoder.load_state_dict( state_dict["autoencoder"] )
self.text_encoder.load_state_dict( state_dict["text_encoder"] )
self.diffusion.load_state_dict( state_dict["diffusion"] )
# = = = = = load from our ckpt = = = = = #
if self.config.ckpt == "real":
print("Saving all real images...")
self.just_save_real = True
else:
checkpoint = torch.load(self.config.ckpt, map_location="cpu")
which_state = 'ema' if 'ema' in checkpoint else "model"
which_state = which_state if config.which_state is None else config.which_state
self.model.load_state_dict(checkpoint[which_state])
print("ckpt is loaded")
self.just_save_real = False
set_alpha_scale(self.model, self.config.alpha_scale)
self.autoencoder.eval()
self.model.eval()
self.text_encoder.eval()
# = = = = = create data = = = = = #
self.dataset_eval = ConCatDataset(config.val_dataset_names, config.DATA_ROOT, config.which_embedder, train=False)
print("total eval images: ", len(self.dataset_eval))
sampler = DistributedSampler(self.dataset_eval,shuffle=False) if config.distributed else None
loader_eval = DataLoader( self.dataset_eval,batch_size=config.batch_size,
num_workers=config.workers,
pin_memory=True,
sampler=sampler,
drop_last=False) # shuffle default is False
self.loader_eval = loader_eval
# = = = = = create output folder = = = = = #
folder_name = ckpt_to_folder_name(os.path.basename(config.ckpt))
self.outdir = os.path.join(config.OUTPUT_ROOT, folder_name)
self.outdir_real = os.path.join(self.outdir,'real')
self.outdir_fake = os.path.join(self.outdir,'fake')
if config.to256:
self.outdir_real256 = os.path.join(self.outdir,'real256')
self.outdir_fake256 = os.path.join(self.outdir,'fake256')
synchronize() # if rank0 is faster, it may mkdir before the other rank call os.listdir()
if get_rank() == 0:
os.makedirs(self.outdir, exist_ok=True)
os.makedirs(self.outdir_real, exist_ok=True)
os.makedirs(self.outdir_fake, exist_ok=True)
if config.to256:
os.makedirs(self.outdir_real256, exist_ok=True)
os.makedirs(self.outdir_fake256, exist_ok=True)
print(self.outdir) # double check
self.evaluation_finished = False
if os.path.exists( os.path.join(self.outdir,'score.txt') ):
self.evaluation_finished = True
def alread_saved_this_batch(self, batch):
existing_real_files = os.listdir( self.outdir_real )
existing_fake_files = os.listdir( self.outdir_fake )
status = []
for image_id in batch["id"]:
img_name = str(int(image_id))+'.png'
status.append(img_name in existing_real_files)
status.append(img_name in existing_fake_files)
return all(status)
@torch.no_grad()
def start_evaluating(self):
iterator = tqdm( self.loader_eval, desc='Evaluating progress')
for batch in iterator:
#if not self.alread_saved_this_batch(batch):
if True:
batch_to_device(batch, self.device)
batch_size = batch["image"].shape[0]
samples_real = batch["image"]
if self.just_save_real:
samples_fake = None
else:
uc = self.text_encoder.encode( batch_size*[""] )
context = self.text_encoder.encode( batch["caption"] )
image_mask = x0 = None
if self.config.inpaint:
image_mask = draw_masks_from_boxes( batch['boxes'], self.model.image_size ).cuda()
x0 = self.autoencoder.encode( batch["image"] )
shape = (batch_size, self.model.in_channels, self.model.image_size, self.model.image_size)
if self.config.no_plms:
sampler = DDIMSampler(self.diffusion, self.model)
steps = 250
else:
sampler = PLMSSampler(self.diffusion, self.model)
steps = 50
input = dict( x=None, timesteps=None, context=context, boxes=batch['boxes'], masks=batch['masks'], positive_embeddings=batch["positive_embeddings"] )
samples_fake = sampler.sample(S=steps, shape=shape, input=input, uc=uc, guidance_scale=self.config.guidance_scale, mask=image_mask, x0=x0)
samples_fake = self.autoencoder.decode(samples_fake)
save_images(samples_real, batch['id'], self.outdir_real, to256=False )
if self.config.to256:
save_images(samples_real, batch['id'], self.outdir_real256, to256=True )
if samples_fake is not None:
save_images(samples_fake, batch['id'], self.outdir_fake, to256=False )
if self.config.to256:
save_images(samples_fake, batch['id'], self.outdir_fake256, to256=True )
def fire_fid(self):
paths = [self.outdir_real, self.outdir_fake]
if self.config.to256:
paths = [self.outdir_real256, self.outdir_fake256]