evaluation/general_score_pap_val.py

from datasets import load_dataset
from pick_score import PickScorer
from aesthetic_score import AestheticScorer
from hpsv2_score import HPSv2Scorer
from imagereward_score import load_imagereward
from diffusers import AutoencoderKL, StableDiffusionPipeline, \
                        StableDiffusionXLPipeline, DDIMScheduler, \
                        UNet2DConditionModel
import torch
import os
import json
from tqdm import tqdm
from huggingface_hub import hf_hub_download
from argparse import ArgumentParser


def load_origin_sd_v1_5(scheduler, inference_dtype):
    pipe = StableDiffusionPipeline.from_pretrained(
        'stable-diffusion-v1-5/stable-diffusion-v1-5',
        torch_dtype=inference_dtype,
        scheduler=scheduler,
        safety_checker=None,
    )
    guidance_scale = 7.5
    return pipe, guidance_scale


def load_spo_sd_v1_5(scheduler, inference_dtype):
    pipe = StableDiffusionPipeline.from_pretrained(
        'SPO-Diffusion-Models/SPO-SD-v1-5_4k-p_10ep',
        torch_dtype=inference_dtype,
        scheduler=scheduler,
        safety_checker=None,
    )
    guidance_scale = 5.0
    return pipe, guidance_scale


def load_diffusion_dpo_sd_v1_5(scheduler, inference_dtype):
    unet = UNet2DConditionModel.from_pretrained('mhdang/dpo-sd1.5-text2image-v1', subfolder="unet", torch_dtype=inference_dtype)
    pipe = StableDiffusionPipeline.from_pretrained(
        'stable-diffusion-v1-5/stable-diffusion-v1-5',
        torch_dtype=inference_dtype,
        scheduler=scheduler,
        safety_checker=None,
        unet=unet,
    )
    guidance_scale = 7.5
    return pipe, guidance_scale


def load_lpo_sd_v1_5(scheduler, inference_dtype):
    unet = UNet2DConditionModel.from_pretrained(
        'casiatao/LPO',
        subfolder="lpo_sd15_merge/unet",
        torch_dtype=inference_dtype
    )
    pipe = StableDiffusionPipeline.from_pretrained(
        'stable-diffusion-v1-5/stable-diffusion-v1-5',
        torch_dtype=inference_dtype,
        scheduler=scheduler,
        safety_checker=None,
        unet=unet
    )
    guidance_scale = 5.0
    return pipe, guidance_scale


def load_origin_sdxl(scheduler, inference_dtype):
    pipe = StableDiffusionXLPipeline.from_pretrained(
        'stabilityai/stable-diffusion-xl-base-1.0',
        torch_dtype=inference_dtype,
        scheduler=scheduler,
    )
    vae = AutoencoderKL.from_pretrained(
        'madebyollin/sdxl-vae-fp16-fix',
        torch_dtype=torch.float16,
    )
    pipe.vae = vae
    guidance_scale = 5.0
    return pipe, guidance_scale


def load_spo_sdxl(scheduler, inference_dtype):
    pipe = StableDiffusionXLPipeline.from_pretrained(
        'SPO-Diffusion-Models/SPO-SDXL_4k-p_10ep',
        torch_dtype=inference_dtype,
        scheduler=scheduler,
    )
    vae = AutoencoderKL.from_pretrained(
        'madebyollin/sdxl-vae-fp16-fix',
        torch_dtype=torch.float16,
    )
    pipe.vae = vae
    guidance_scale = 5.0
    return pipe, guidance_scale


def load_diffusion_dpo_sdxl(scheduler, inference_dtype):
    unet = UNet2DConditionModel.from_pretrained('mhdang/dpo-sdxl-text2image-v1', subfolder="unet", torch_dtype=inference_dtype)
    pipe = StableDiffusionXLPipeline.from_pretrained(
        'stabilityai/stable-diffusion-xl-base-1.0',
        torch_dtype=inference_dtype,
        scheduler=scheduler,
        unet=unet,
    )
    vae = AutoencoderKL.from_pretrained(
        'madebyollin/sdxl-vae-fp16-fix',
        torch_dtype=torch.float16,
    )
    pipe.vae = vae
    guidance_scale = 5.0
    return pipe, guidance_scale


def load_lpo_sdxl(scheduler, inference_dtype):
    unet = UNet2DConditionModel.from_pretrained(
        'casiatao/LPO',
        subfolder="lpo_sdxl_merge/unet",
        torch_dtype=inference_dtype
    )
    vae = AutoencoderKL.from_pretrained(
        'madebyollin/sdxl-vae-fp16-fix',
        torch_dtype=torch.float16,
    )
    pipe = StableDiffusionXLPipeline.from_pretrained(
        'stabilityai/stable-diffusion-xl-base-1.0',
        torch_dtype=inference_dtype,
        scheduler=scheduler,
        unet=unet,
        vae=vae
    )
    guidance_scale = 5.0
    return pipe, guidance_scale



model_dict = {
    'origin_sd15': load_origin_sd_v1_5,
    'spo_sd15': load_spo_sd_v1_5,
    'diffusion_dpo_sd15': load_diffusion_dpo_sd_v1_5,
    'lpo_sd15': load_lpo_sd_v1_5,
    'origin_sdxl': load_origin_sdxl,
    'spo_sdxl': load_spo_sdxl,
    'diffusion_dpo_sdxl': load_diffusion_dpo_sdxl,
    'lpo_sdxl': load_lpo_sdxl,
}


if __name__ == "__main__":
    # hyperparameter
    parser = ArgumentParser()
    parser.add_argument("--model_name", type=str, default="origin_sdxl")
    parser.add_argument("--batch_size", type=int, default=1)
    parser.add_argument("--num_image_per_prompt", type=int, default=4)
    parser.add_argument("--sample_steps", type=int, default=20)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--device", type=str, default="cuda")
    args = parser.parse_args()

    model_name = args.model_name
    batch_size = args.batch_size
    num_image_per_prompt = args.num_image_per_prompt
    sample_steps = args.sample_steps
    seed = args.seed
    device = args.device

    # load preference model
    pickscorer = PickScorer(processor_name_or_path="laion/CLIP-ViT-H-14-laion2B-s32B-b79K", model_pretrained_name_or_path="yuvalkirstain/PickScore_v1", device=device)
    
    aesthetic_scorer = AestheticScorer(torch.float32, "openai/clip-vit-large-patch14", "./sac+logos+ava1-l14-linearMSE.pth")
    aesthetic_scorer = aesthetic_scorer.to(device)
    

    hpsv2_scorer = HPSv2Scorer(
        clip_pretrained_name_or_path=hf_hub_download(repo_id="laion/CLIP-ViT-H-14-laion2B-s32B-b79K", filename="open_clip_pytorch_model.bin"),
        model_pretrained_name_or_path=hf_hub_download(repo_id="xswu/HPSv2", filename="HPS_v2_compressed.pt"),
        device=device
    )

    hpsv21_scorer = HPSv2Scorer(
        clip_pretrained_name_or_path=hf_hub_download(repo_id="laion/CLIP-ViT-H-14-laion2B-s32B-b79K", filename="open_clip_pytorch_model.bin"),
        model_pretrained_name_or_path=hf_hub_download(repo_id="xswu/HPSv2", filename="HPS_v2.1_compressed.pt"),
        device=device
    )

    imagereward_scorer = load_imagereward(
        model_path=hf_hub_download(repo_id="THUDM/ImageReward", filename="ImageReward.pt"), 
        med_config=hf_hub_download(repo_id="THUDM/ImageReward", filename="med_config.json"), 
        device=device
    )
    
    load_model_func = model_dict[model_name]

    # load diffusion model
    inference_dtype = torch.float16
    scheduler = DDIMScheduler.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="scheduler")
    pipe, guidance_scale = load_model_func(scheduler, inference_dtype)
    pipe.to(device)
    pipe.set_progress_bar_config(disable=True)
    
    # load dataset
    val_dataset = load_dataset("pickapic-anonymous/pickapic_v1", split="validation_unique", streaming=True)

    # calculate preference score
    caption_list = []
    for i, sample in enumerate(val_dataset):
        caption_list.append(sample['caption'])
    
    batch_num = len(caption_list) // batch_size if len(caption_list) % batch_size == 0 else len(caption_list) // batch_size + 1
    batched_caption_list = [caption_list[i*batch_size:(i+1)*batch_size] for i in range(batch_num)]
    
    pickscore_list = []
    aesthetic_score_list = []
    hpsv2score_list = []
    hpsv21score_list = []
    imagereward_list = []
    
    for batch_prompt in tqdm(batched_caption_list):
        generator=torch.Generator(device=device).manual_seed(seed)
        images = pipe(
            batch_prompt,
            guidance_scale=guidance_scale,
            num_inference_steps=sample_steps,
            generator=generator,
            output_type='pil',
            num_images_per_prompt=num_image_per_prompt,
        ).images
        
        for prompt, image in zip(batch_prompt, images):
            pickscore = pickscorer(prompt, [image])[0]
            pickscore_list.append(pickscore)
            
            aesthetic_score = aesthetic_scorer(image)[0].item()
            aesthetic_score_list.append(aesthetic_score)
            
            hpsv2_score = hpsv2_scorer.score(image, prompt)[0]
            hpsv2score_list.append(hpsv2_score)
            
            hpsv21_score = hpsv21_scorer.score(image, prompt)[0]
            hpsv21score_list.append(hpsv21_score)
            
            imagereward_score = imagereward_scorer.score(prompt, image)
            imagereward_list.append(imagereward_score)
    
    
    res_save_dir = './eval_results/pick_a_pic_val_score'
    os.makedirs(res_save_dir, exist_ok=True)

    file_name = f"{model_name}_ddim_cfg{guidance_scale}_step{sample_steps}_seed{seed}_{num_image_per_prompt}image_batch{batch_size}.json"
    with open(os.path.join(res_save_dir, file_name), 'w', encoding='utf-8') as f:
        json.dump({
            'pickscore': torch.mean(torch.tensor(pickscore_list)).item(),
            'aestheticscore': torch.mean(torch.tensor(aesthetic_score_list)).item(),
            'hpsv2score': torch.mean(torch.tensor(hpsv2score_list)).item(),
            'hpsv21score': torch.mean(torch.tensor(hpsv21score_list)).item(),
            'imagerewardscore': torch.mean(torch.tensor(imagereward_list)).item(),
        }, f, indent=4)
    
    print(f"Pickscore: {torch.mean(torch.tensor(pickscore_list))}")
    print(f"Aestheticscore: {torch.mean(torch.tensor(aesthetic_score_list))}")
    print(f"HPSv2score: {torch.mean(torch.tensor(hpsv2score_list))}")
    print(f"HPSv21score: {torch.mean(torch.tensor(hpsv21score_list))}")
    print(f"Imagerewardscore: {torch.mean(torch.tensor(imagereward_list))}")