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 #hello 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]