Upload 24 files

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/concept.png filter=lfs diff=lfs merge=lfs -text
+assets/examples.png filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 EML
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

arguments.py

@@ -0,0 +1,111 @@
+import argparse
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Process Reward Optimization.")
+
+    # update paths here!
+    parser.add_argument(
+        "--cache_dir",
+        type=str,
+        help="HF cache directory",
+        default="/shared-local/aoq951/HF_CACHE/",
+    )
+    parser.add_argument(
+        "--save_dir",
+        type=str,
+        help="Directory to save images",
+        default="/shared-local/aoq951/ReNO/outputs",
+    )
+
+    # model and optim
+    parser.add_argument("--model", type=str, help="Model to use", default="sdxl-turbo")
+    parser.add_argument("--lr", type=float, help="Learning rate", default=5.0)
+    parser.add_argument("--n_iters", type=int, help="Number of iterations", default=50)
+    parser.add_argument(
+        "--n_inference_steps", type=int, help="Number of iterations", default=1
+    )
+    parser.add_argument(
+        "--optim",
+        choices=["sgd", "adam", "lbfgs"],
+        default="sgd",
+        help="Optimizer to be used",
+    )
+    parser.add_argument("--nesterov", default=True, action="store_false")
+    parser.add_argument(
+        "--grad_clip", type=float, help="Gradient clipping", default=0.1
+    )
+    parser.add_argument("--seed", type=int, help="Seed to use", default=0)
+
+    # reward losses
+    parser.add_argument("--disable_hps", default=True, action="store_false",dest="enable_hps")
+    parser.add_argument(
+        "--hps_weighting", type=float, help="Weighting for HPS", default=5.0
+    )
+    parser.add_argument("--disable_imagereward", default=True, action="store_false",dest='enable_imagereward')
+    parser.add_argument(
+        "--imagereward_weighting",
+        type=float,
+        help="Weighting for ImageReward",
+        default=1.0,
+    )
+    parser.add_argument("--disable_clip", default=True, action="store_false",dest='enable_clip')
+    parser.add_argument(
+        "--clip_weighting", type=float, help="Weighting for CLIP", default=0.01
+    )
+    parser.add_argument("--disable_pickscore", default=True, action="store_false",dest='enable_pickscore')
+    parser.add_argument(
+        "--pickscore_weighting",
+        type=float,
+        help="Weighting for PickScore",
+        default=0.05,
+    )
+    parser.add_argument("--disable_aesthetic", default=False, action="store_false",dest='enable_aesthetic')
+    parser.add_argument(
+        "--aesthetic_weighting",
+        type=float,
+        help="Weighting for Aesthetic",
+        default=0.0,
+    )
+    parser.add_argument("--disable_reg", default=True, action="store_false",dest='enable_reg')
+    parser.add_argument(
+        "--reg_weight", type=float, help="Regularization weight", default=0.01
+    )
+
+    # task specific
+    parser.add_argument(
+        "--task",
+        type=str,
+        help="Task to run",
+        default="single",
+        choices=[
+            "t2i-compbench",
+            "single",
+            "parti-prompts",
+            "geneval",
+            "example-prompts",
+        ],
+    )
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        help="Prompt to run",
+        default="A red dog and a green cat",
+    )
+    parser.add_argument(
+        "--benchmark_reward",
+        help="Reward to benchmark on",
+        default="total",
+        choices=["ImageReward", "PickScore", "HPS", "CLIP", "total"],
+    )
+
+    # general
+    parser.add_argument("--save_all_images", default=False, action="store_true")
+    parser.add_argument("--no_optim", default=False, action="store_true")
+    parser.add_argument("--imageselect", default=False, action="store_true")
+    parser.add_argument("--memsave", default=False, action="store_true")
+    parser.add_argument("--device", type=str, help="Device to use", default="cuda")
+    parser.add_argument("--device_id", type=int, help="Device ID to use", default=None)
+
+    args = parser.parse_args()
+    return args

assets/example_prompts.txt

@@ -0,0 +1,30 @@
+A minimalist logo design of a reindeer, fully rendered. The reindeer features distinct, complete shapes using bold and flat colors. The design emphasizes simplicity and clarity, suitable for logo use with a sharp outline and white background.
+A red dog and a green cat
+A green dog and a red cat
+A pink elephant and a grey cow
+A grey elephant and a pink cow
+A yellow reindeer and a blue elephant
+A blue reindeer and a yellow elephant
+An orange chair to the right of a black airplane
+Three dogs and two horses
+A cat playing checkers
+High quality photo of a monkey astronaut infront of the Eiffel tower
+A bird with 8 legs
+A brain riding a rocketship towards the moon
+A toaster riding a bike
+A blue scooter is parked near a curb in front of a green vintage car
+A curious, orange fox and a fluffy, white rabbit, playing together in a lush, green meadow filled with yellow dandelions
+An epic oil painting: a red portal infront of a cityscape, a solitary figure, and a colorful sky over snowy mountains
+A futuristic painting: Red car escapes giant shark's leap, right; ominous mountains, blue sky
+A majestic, resilient sea ship navigates the icy wilderness in the style of Star Wars
+Dwayne Johnson depicted as a philosopher king in an academic painting by Greg Rutkowski
+Taylor Swift depicted as a prime minister in an academic painting by Kandinsky
+A watercolor painting: a floating island, multiple animals under a majestic tree with golden leaves, and a vibrant rainbow stretching across a pastel sky
+A Japanese-style ink painting: a traditional wooden bridge, a pagoda, a lone samurai warrior, and cherry blossom petals over a tranquil river
+A retro-futuristic pixel art scene: a flying car, an imperial senate building, a green park, and a purple sunset
+A impressionistic oil painting: a lone figure walking on a misty beach, a weathered lighthouse on a cliff, seagulls above crashing waves
+A fairytale castle with a golden-haired woman in a floral-patterned metallic frame, and a vase with lilies
+A post-apocalyptic digital artwork: crumbling skyscrapers, an abandoned car overgrown with vines, and a fiery orange sunset casting long shadows
+A stop sign infront of a traffic light
+A sign on a grocery store that has 'ENTRY' written on it
+A colorful poster with the title 'INTERGALACTICAL' written on it

environment.yml

@@ -0,0 +1,24 @@
+name: reno
+channels:
+  - pytorch
+  - nvidia
+  - conda-forge
+dependencies:
+  - python=3.11
+  - pytorch=2.3
+  - torchvision=0.18.0
+  - pytorch-cuda
+  - pytorch-lightning=2.2
+  - pip
+  - pip:
+    - datasets==2.18
+    - transformers==4.38.2
+    - diffusers==0.28
+    - hpsv2==1.2
+    - image-reward==1.5
+    - open-clip-torch==2.24
+    - blobfile
+    - openai-clip
+    - setuptools==60.2
+    - optimum
+    - xformers

main.py

@@ -0,0 +1,274 @@
+import json
+import logging
+import os
+
+import blobfile as bf
+import torch
+from datasets import load_dataset
+from pytorch_lightning import seed_everything
+from tqdm import tqdm
+
+from arguments import parse_args
+from models import get_model
+from rewards import get_reward_losses
+from training import LatentNoiseTrainer, get_optimizer
+
+
+def main(args):
+    seed_everything(args.seed)
+    bf.makedirs(f"{args.save_dir}/logs/{args.task}")
+    # Set up logging and name settings
+    logger = logging.getLogger()
+    settings = (
+        f"{args.model}{'_' + args.prompt if args.task == 't2i-compbench' else ''}"
+        f"{'_no-optim' if args.no_optim else ''}_{args.seed if args.task != 'geneval' else ''}"
+        f"_lr{args.lr}_gc{args.grad_clip}_iter{args.n_iters}"
+        f"_reg{args.reg_weight if args.enable_reg else '0'}"
+        f"{'_pickscore' + str(args.pickscore_weighting) if args.enable_pickscore else ''}"
+        f"{'_clip' + str(args.clip_weighting) if args.enable_clip else ''}"
+        f"{'_hps' + str(args.hps_weighting) if args.enable_hps else ''}"
+        f"{'_imagereward' + str(args.imagereward_weighting) if args.enable_imagereward else ''}"
+        f"{'_aesthetic' + str(args.aesthetic_weighting) if args.enable_aesthetic else ''}"
+    )
+    file_stream = open(f"{args.save_dir}/logs/{args.task}/{settings}.txt", "w")
+    handler = logging.StreamHandler(file_stream)
+    formatter = logging.Formatter("%(asctime)s - %(message)s")
+    handler.setFormatter(formatter)
+    logger.addHandler(handler)
+    logger.setLevel("INFO")
+    consoleHandler = logging.StreamHandler()
+    consoleHandler.setFormatter(formatter)
+    logger.addHandler(consoleHandler)
+    logging.info(args)
+    if args.device_id is not None:
+        logging.info(f"Using CUDA device {args.device_id}")
+        os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
+        os.environ["CUDA_VISIBLE_DEVICE"] = args.device_id
+    if args.device == "cuda":
+        device = torch.device("cuda")
+    else:
+        device = torch.device("cpu")
+    # Set dtype to fp16
+    dtype = torch.float16
+    # Get reward losses
+    reward_losses = get_reward_losses(args, dtype, device, args.cache_dir)
+
+    # Get model and noise trainer
+    sd_model = get_model(args.model, dtype, device, args.cache_dir, args.memsave)
+    trainer = LatentNoiseTrainer(
+        reward_losses=reward_losses,
+        model=sd_model,
+        n_iters=args.n_iters,
+        n_inference_steps=args.n_inference_steps,
+        seed=args.seed,
+        save_all_images=args.save_all_images,
+        device=device,
+        no_optim=args.no_optim,
+        regularize=args.enable_reg,
+        regularization_weight=args.reg_weight,
+        grad_clip=args.grad_clip,
+        log_metrics=args.task == "single" or not args.no_optim,
+        imageselect=args.imageselect,
+    )
+
+    # Create latents
+    if args.model != "pixart":
+        height = sd_model.unet.config.sample_size * sd_model.vae_scale_factor
+        width = sd_model.unet.config.sample_size * sd_model.vae_scale_factor
+        shape = (
+            1,
+            sd_model.unet.in_channels,
+            height // sd_model.vae_scale_factor,
+            width // sd_model.vae_scale_factor,
+        )
+    else:
+        height = sd_model.transformer.config.sample_size * sd_model.vae_scale_factor
+        width = sd_model.transformer.config.sample_size * sd_model.vae_scale_factor
+        shape = (
+            1,
+            sd_model.transformer.config.in_channels,
+            height // sd_model.vae_scale_factor,
+            width // sd_model.vae_scale_factor,
+        )
+    enable_grad = not args.no_optim
+
+    if args.task == "single":
+        init_latents = torch.randn(shape, device=device, dtype=dtype)
+        latents = torch.nn.Parameter(init_latents, requires_grad=enable_grad)
+        optimizer = get_optimizer(args.optim, latents, args.lr, args.nesterov)
+        save_dir = f"{args.save_dir}/{args.task}/{settings}/{args.prompt}"
+        os.makedirs(f"{save_dir}", exist_ok=True)
+        best_image, total_init_rewards, total_best_rewards = trainer.train(
+            latents, args.prompt, optimizer, save_dir
+        )
+        best_image.save(f"{save_dir}/best_image.png")
+    elif args.task == "example-prompts":
+        fo = open("assets/example_prompts.txt", "r")
+        prompts = fo.readlines()
+        fo.close()
+        for i, prompt in tqdm(enumerate(prompts)):
+            # Get new latents and optimizer
+            init_latents = torch.randn(shape, device=device, dtype=dtype)
+            latents = torch.nn.Parameter(init_latents, requires_grad=enable_grad)
+            optimizer = get_optimizer(args.optim, latents, args.lr, args.nesterov)
+
+            prompt = prompt.strip()
+            name = f"{i:03d}_{prompt}.png"
+            save_dir = f"{args.save_dir}/{args.task}/{settings}/{name}"
+            os.makedirs(save_dir, exist_ok=True)
+            best_image, init_rewards, best_rewards = trainer.train(
+                latents, prompt, optimizer, save_dir
+            )
+            if i == 0:
+                total_best_rewards = {k: 0.0 for k in best_rewards.keys()}
+                total_init_rewards = {k: 0.0 for k in best_rewards.keys()}
+            for k in best_rewards.keys():
+                total_best_rewards[k] += best_rewards[k]
+                total_init_rewards[k] += init_rewards[k]
+            best_image.save(f"{save_dir}/best_image.png")
+            logging.info(f"Initial rewards: {init_rewards}")
+            logging.info(f"Best rewards: {best_rewards}")
+        for k in total_best_rewards.keys():
+            total_best_rewards[k] /= len(prompts)
+            total_init_rewards[k] /= len(prompts)
+
+        # save results to directory
+        with open(f"{args.save_dir}/example-prompts/{settings}/results.txt", "w") as f:
+            f.write(
+                f"Mean initial all rewards: {total_init_rewards}\n"
+                f"Mean best all rewards: {total_best_rewards}\n"
+            )
+    elif args.task == "t2i-compbench":
+        prompt_list_file = f"../T2I-CompBench/examples/dataset/{args.prompt}.txt"
+        fo = open(prompt_list_file, "r")
+        prompts = fo.readlines()
+        fo.close()
+        os.makedirs(f"{args.save_dir}/{args.task}/{settings}/samples", exist_ok=True)
+        for i, prompt in tqdm(enumerate(prompts)):
+            # Get new latents and optimizer
+            init_latents = torch.randn(shape, device=device, dtype=dtype)
+            latents = torch.nn.Parameter(init_latents, requires_grad=enable_grad)
+            optimizer = get_optimizer(args.optim, latents, args.lr, args.nesterov)
+
+            prompt = prompt.strip()
+            best_image, init_rewards, best_rewards = trainer.train(
+                latents, prompt, optimizer
+            )
+            if i == 0:
+                total_best_rewards = {k: 0.0 for k in best_rewards.keys()}
+                total_init_rewards = {k: 0.0 for k in best_rewards.keys()}
+            for k in best_rewards.keys():
+                total_best_rewards[k] += best_rewards[k]
+                total_init_rewards[k] += init_rewards[k]
+            name = f"{prompt}_{i:06d}.png"
+            best_image.save(f"{args.save_dir}/{args.task}/{settings}/samples/{name}")
+            logging.info(f"Initial rewards: {init_rewards}")
+            logging.info(f"Best rewards: {best_rewards}")
+        for k in total_best_rewards.keys():
+            total_best_rewards[k] /= len(prompts)
+            total_init_rewards[k] /= len(prompts)
+    elif args.task == "parti-prompts":
+        parti_dataset = load_dataset("nateraw/parti-prompts", split="train")
+        total_reward_diff = 0.0
+        total_best_reward = 0.0
+        total_init_reward = 0.0
+        total_improved_samples = 0
+        for index, sample in enumerate(parti_dataset):
+            os.makedirs(
+                f"{args.save_dir}/{args.task}/{settings}/{index}", exist_ok=True
+            )
+            prompt = sample["Prompt"]
+            best_image, init_rewards, best_rewards = trainer.train(
+                latents, prompt, optimizer
+            )
+            best_image.save(
+                f"{args.save_dir}/{args.task}/{settings}/{index}/best_image.png"
+            )
+            open(
+                f"{args.save_dir}/{args.task}/{settings}/{index}/prompt.txt", "w"
+            ).write(
+                f"{prompt} \n Initial Rewards: {init_rewards} \n Best Rewards: {best_rewards}"
+            )
+            logging.info(f"Initial rewards: {init_rewards}")
+            logging.info(f"Best rewards: {best_rewards}")
+            initial_reward = init_rewards[args.benchmark_reward]
+            best_reward = best_rewards[args.benchmark_reward]
+            total_reward_diff += best_reward - initial_reward
+            total_best_reward += best_reward
+            total_init_reward += initial_reward
+            if best_reward < initial_reward:
+                total_improved_samples += 1
+            if i == 0:
+                total_best_rewards = {k: 0.0 for k in best_rewards.keys()}
+                total_init_rewards = {k: 0.0 for k in best_rewards.keys()}
+            for k in best_rewards.keys():
+                total_best_rewards[k] += best_rewards[k]
+                total_init_rewards[k] += init_rewards[k]
+            # Get new latents and optimizer
+            init_latents = torch.randn(shape, device=device, dtype=dtype)
+            latents = torch.nn.Parameter(init_latents, requires_grad=enable_grad)
+            optimizer = get_optimizer(args.optim, latents, args.lr, args.nesterov)
+        improvement_percentage = total_improved_samples / parti_dataset.num_rows
+        mean_best_reward = total_best_reward / parti_dataset.num_rows
+        mean_init_reward = total_init_reward / parti_dataset.num_rows
+        mean_reward_diff = total_reward_diff / parti_dataset.num_rows
+        logging.info(
+            f"Improvement percentage: {improvement_percentage:.4f}, "
+            f"mean initial reward: {mean_init_reward:.4f}, "
+            f"mean best reward: {mean_best_reward:.4f}, "
+            f"mean reward diff: {mean_reward_diff:.4f}"
+        )
+        for k in total_best_rewards.keys():
+            total_best_rewards[k] /= len(parti_dataset)
+            total_init_rewards[k] /= len(parti_dataset)
+        # save results
+        os.makedirs(f"{args.save_dir}/parti-prompts/{settings}", exist_ok=True)
+        with open(f"{args.save_dir}/parti-prompts/{settings}/results.txt", "w") as f:
+            f.write(
+                f"Mean improvement: {improvement_percentage:.4f}, "
+                f"mean initial reward: {mean_init_reward:.4f}, "
+                f"mean best reward: {mean_best_reward:.4f}, "
+                f"mean reward diff: {mean_reward_diff:.4f}\n"
+                f"Mean initial all rewards: {total_init_rewards}\n"
+                f"Mean best all rewards: {total_best_rewards}"
+            )
+    elif args.task == "geneval":
+        prompt_list_file = "../geneval/prompts/evaluation_metadata.jsonl"
+        with open(prompt_list_file) as fp:
+            metadatas = [json.loads(line) for line in fp]
+        outdir = f"{args.save_dir}/{args.task}/{settings}"
+        for index, metadata in enumerate(metadatas):
+            # Get new latents and optimizer
+            init_latents = torch.randn(shape, device=device, dtype=dtype)
+            latents = torch.nn.Parameter(init_latents, requires_grad=True)
+            optimizer = get_optimizer(args.optim, latents, args.lr, args.nesterov)
+
+            prompt = metadata["prompt"]
+            best_image, init_rewards, best_rewards = trainer.train(
+                latents, prompt, optimizer
+            )
+            logging.info(f"Initial rewards: {init_rewards}")
+            logging.info(f"Best rewards: {best_rewards}")
+            outpath = f"{outdir}/{index:0>5}"
+            os.makedirs(f"{outpath}/samples", exist_ok=True)
+            with open(f"{outpath}/metadata.jsonl", "w") as fp:
+                json.dump(metadata, fp)
+            best_image.save(f"{outpath}/samples/{args.seed:05}.png")
+            if i == 0:
+                total_best_rewards = {k: 0.0 for k in best_rewards.keys()}
+                total_init_rewards = {k: 0.0 for k in best_rewards.keys()}
+            for k in best_rewards.keys():
+                total_best_rewards[k] += best_rewards[k]
+                total_init_rewards[k] += init_rewards[k]
+        for k in total_best_rewards.keys():
+            total_best_rewards[k] /= len(parti_dataset)
+            total_init_rewards[k] /= len(parti_dataset)
+    else:
+        raise ValueError(f"Unknown task {args.task}")
+    # log total rewards
+    logging.info(f"Mean initial rewards: {total_init_rewards}")
+    logging.info(f"Mean best rewards: {total_best_rewards}")
+    
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)

models/RewardPixart.py

@@ -0,0 +1,394 @@
+from typing import List, Optional, Union
+
+import torch
+from diffusers import PixArtAlphaPipeline
+from diffusers.pipelines.pixart_alpha.pipeline_pixart_alpha import \
+    retrieve_timesteps
+
+
+def freeze_params(params):
+    for param in params:
+        param.requires_grad = False
+
+
+class RewardPixartPipeline(PixArtAlphaPipeline):
+    def __init__(
+        self, tokenizer, text_encoder, transformer, scheduler, vae, memsave=False
+    ):
+        super().__init__(
+            tokenizer,
+            text_encoder,
+            vae,
+            transformer,
+            scheduler,
+        )
+        # optionally enable memsave_torch
+        if memsave:
+            import memsave_torch.nn
+
+            self.vae = memsave_torch.nn.convert_to_memory_saving(self.vae)
+            self.text_encoder = memsave_torch.nn.convert_to_memory_saving(
+                self.text_encoder
+            )
+        self.text_encoder.gradient_checkpointing_enable()
+        self.vae.enable_gradient_checkpointing()
+        self.text_encoder.eval()
+        self.vae.eval()
+        freeze_params(self.vae.parameters())
+        freeze_params(self.text_encoder.parameters())
+
+    def apply(
+        self,
+        latents: torch.Tensor = None,
+        prompt: Union[str, List[str]] = None,
+        negative_prompt: str = "",
+        num_inference_steps: int = 20,
+        timesteps: List[int] = [400],
+        sigmas: List[float] = None,
+        guidance_scale: float = 1.0,
+        num_images_per_prompt: Optional[int] = 1,
+        height: Optional[int] = 512,
+        width: Optional[int] = 512,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        prompt_attention_mask: Optional[torch.FloatTensor] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_prompt_attention_mask: Optional[torch.FloatTensor] = None,
+        callback_steps: int = 1,
+        clean_caption: bool = False,
+        use_resolution_binning: bool = True,
+        max_sequence_length: int = 120,
+        **kwargs,
+    ):
+        # 1. Check inputs. Raise error if not correct
+        height = height or self.transformer.config.sample_size * self.vae_scale_factor
+        width = width or self.transformer.config.sample_size * self.vae_scale_factor
+        if use_resolution_binning:
+            if self.transformer.config.sample_size == 128:
+                aspect_ratio_bin = ASPECT_RATIO_1024_BIN
+            elif self.transformer.config.sample_size == 64:
+                aspect_ratio_bin = ASPECT_RATIO_512_BIN
+            elif self.transformer.config.sample_size == 32:
+                aspect_ratio_bin = ASPECT_RATIO_256_BIN
+            else:
+                raise ValueError("Invalid sample size")
+            orig_height, orig_width = height, width
+            height, width = self.image_processor.classify_height_width_bin(
+                height, width, ratios=aspect_ratio_bin
+            )
+
+        self.check_inputs(
+            prompt,
+            height,
+            width,
+            negative_prompt,
+            callback_steps,
+            prompt_embeds,
+            negative_prompt_embeds,
+            prompt_attention_mask,
+            negative_prompt_attention_mask,
+        )
+
+        # 2. Default height and width to transformer
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+
+        # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+        # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+        # corresponds to doing no classifier free guidance.
+        do_classifier_free_guidance = guidance_scale > 1.0
+
+        # 3. Encode input prompt
+        (
+            prompt_embeds,
+            prompt_attention_mask,
+            negative_prompt_embeds,
+            negative_prompt_attention_mask,
+        ) = self.encode_prompt(
+            prompt,
+            do_classifier_free_guidance,
+            negative_prompt=negative_prompt,
+            num_images_per_prompt=num_images_per_prompt,
+            device=device,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            prompt_attention_mask=prompt_attention_mask,
+            negative_prompt_attention_mask=negative_prompt_attention_mask,
+            clean_caption=clean_caption,
+            max_sequence_length=max_sequence_length,
+        )
+        if do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+            prompt_attention_mask = torch.cat(
+                [negative_prompt_attention_mask, prompt_attention_mask], dim=0
+            )
+
+        # 4. Prepare timesteps
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler, num_inference_steps, device, timesteps, sigmas
+        )
+
+        # 5. Prepare latents.
+        latent_channels = self.transformer.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            latent_channels,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 6.1 Prepare micro-conditions.
+        added_cond_kwargs = {"resolution": None, "aspect_ratio": None}
+        if self.transformer.config.sample_size == 128:
+            resolution = torch.tensor([height, width]).repeat(
+                batch_size * num_images_per_prompt, 1
+            )
+            aspect_ratio = torch.tensor([float(height / width)]).repeat(
+                batch_size * num_images_per_prompt, 1
+            )
+            resolution = resolution.to(dtype=prompt_embeds.dtype, device=device)
+            aspect_ratio = aspect_ratio.to(dtype=prompt_embeds.dtype, device=device)
+
+            if do_classifier_free_guidance:
+                resolution = torch.cat([resolution, resolution], dim=0)
+                aspect_ratio = torch.cat([aspect_ratio, aspect_ratio], dim=0)
+
+            added_cond_kwargs = {"resolution": resolution, "aspect_ratio": aspect_ratio}
+
+        # 7. Denoising loop
+        num_warmup_steps = max(
+            len(timesteps) - num_inference_steps * self.scheduler.order, 0
+        )
+
+        for i, t in enumerate(timesteps):
+            latent_model_input = (
+                torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+            )
+            latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+            current_timestep = t
+            if not torch.is_tensor(current_timestep):
+                # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
+                # This would be a good case for the `match` statement (Python 3.10+)
+                is_mps = latent_model_input.device.type == "mps"
+                if isinstance(current_timestep, float):
+                    dtype = torch.float32 if is_mps else torch.float64
+                else:
+                    dtype = torch.int32 if is_mps else torch.int64
+                current_timestep = torch.tensor(
+                    [current_timestep], dtype=dtype, device=latent_model_input.device
+                )
+            elif len(current_timestep.shape) == 0:
+                current_timestep = current_timestep[None].to(latent_model_input.device)
+            # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+            current_timestep = current_timestep.expand(latent_model_input.shape[0])
+
+            # predict noise model_output
+            noise_pred = self.transformer(
+                latent_model_input,
+                encoder_hidden_states=prompt_embeds,
+                encoder_attention_mask=prompt_attention_mask,
+                timestep=current_timestep,
+                added_cond_kwargs=added_cond_kwargs,
+                return_dict=False,
+            )[0]
+
+            # perform guidance
+            if do_classifier_free_guidance:
+                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                noise_pred = noise_pred_uncond + guidance_scale * (
+                    noise_pred_text - noise_pred_uncond
+                )
+
+            # learned sigma
+            if self.transformer.config.out_channels // 2 == latent_channels:
+                noise_pred = noise_pred.chunk(2, dim=1)[0]
+            else:
+                noise_pred = noise_pred
+
+            # compute previous image: x_t -> x_t-1
+            if num_inference_steps == 1:
+                # For DMD one step sampling: https://arxiv.org/abs/2311.18828
+                latents = self.scheduler.step(
+                    noise_pred, t, latents, **extra_step_kwargs
+                ).pred_original_sample
+
+        image = self.vae.decode(
+            latents / self.vae.config.scaling_factor, return_dict=False
+        )[0]
+        if use_resolution_binning:
+            image = self.image_processor.resize_and_crop_tensor(
+                image, orig_width, orig_height
+            )
+
+        image = (image / 2 + 0.5).clamp(0, 1)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+        return image
+
+
+ASPECT_RATIO_2048_BIN = {
+    "0.25": [1024.0, 4096.0],
+    "0.26": [1024.0, 3968.0],
+    "0.27": [1024.0, 3840.0],
+    "0.28": [1024.0, 3712.0],
+    "0.32": [1152.0, 3584.0],
+    "0.33": [1152.0, 3456.0],
+    "0.35": [1152.0, 3328.0],
+    "0.4": [1280.0, 3200.0],
+    "0.42": [1280.0, 3072.0],
+    "0.48": [1408.0, 2944.0],
+    "0.5": [1408.0, 2816.0],
+    "0.52": [1408.0, 2688.0],
+    "0.57": [1536.0, 2688.0],
+    "0.6": [1536.0, 2560.0],
+    "0.68": [1664.0, 2432.0],
+    "0.72": [1664.0, 2304.0],
+    "0.78": [1792.0, 2304.0],
+    "0.82": [1792.0, 2176.0],
+    "0.88": [1920.0, 2176.0],
+    "0.94": [1920.0, 2048.0],
+    "1.0": [2048.0, 2048.0],
+    "1.07": [2048.0, 1920.0],
+    "1.13": [2176.0, 1920.0],
+    "1.21": [2176.0, 1792.0],
+    "1.29": [2304.0, 1792.0],
+    "1.38": [2304.0, 1664.0],
+    "1.46": [2432.0, 1664.0],
+    "1.67": [2560.0, 1536.0],
+    "1.75": [2688.0, 1536.0],
+    "2.0": [2816.0, 1408.0],
+    "2.09": [2944.0, 1408.0],
+    "2.4": [3072.0, 1280.0],
+    "2.5": [3200.0, 1280.0],
+    "2.89": [3328.0, 1152.0],
+    "3.0": [3456.0, 1152.0],
+    "3.11": [3584.0, 1152.0],
+    "3.62": [3712.0, 1024.0],
+    "3.75": [3840.0, 1024.0],
+    "3.88": [3968.0, 1024.0],
+    "4.0": [4096.0, 1024.0],
+}
+
+ASPECT_RATIO_256_BIN = {
+    "0.25": [128.0, 512.0],
+    "0.28": [128.0, 464.0],
+    "0.32": [144.0, 448.0],
+    "0.33": [144.0, 432.0],
+    "0.35": [144.0, 416.0],
+    "0.4": [160.0, 400.0],
+    "0.42": [160.0, 384.0],
+    "0.48": [176.0, 368.0],
+    "0.5": [176.0, 352.0],
+    "0.52": [176.0, 336.0],
+    "0.57": [192.0, 336.0],
+    "0.6": [192.0, 320.0],
+    "0.68": [208.0, 304.0],
+    "0.72": [208.0, 288.0],
+    "0.78": [224.0, 288.0],
+    "0.82": [224.0, 272.0],
+    "0.88": [240.0, 272.0],
+    "0.94": [240.0, 256.0],
+    "1.0": [256.0, 256.0],
+    "1.07": [256.0, 240.0],
+    "1.13": [272.0, 240.0],
+    "1.21": [272.0, 224.0],
+    "1.29": [288.0, 224.0],
+    "1.38": [288.0, 208.0],
+    "1.46": [304.0, 208.0],
+    "1.67": [320.0, 192.0],
+    "1.75": [336.0, 192.0],
+    "2.0": [352.0, 176.0],
+    "2.09": [368.0, 176.0],
+    "2.4": [384.0, 160.0],
+    "2.5": [400.0, 160.0],
+    "3.0": [432.0, 144.0],
+    "4.0": [512.0, 128.0],
+}
+
+ASPECT_RATIO_1024_BIN = {
+    "0.25": [512.0, 2048.0],
+    "0.28": [512.0, 1856.0],
+    "0.32": [576.0, 1792.0],
+    "0.33": [576.0, 1728.0],
+    "0.35": [576.0, 1664.0],
+    "0.4": [640.0, 1600.0],
+    "0.42": [640.0, 1536.0],
+    "0.48": [704.0, 1472.0],
+    "0.5": [704.0, 1408.0],
+    "0.52": [704.0, 1344.0],
+    "0.57": [768.0, 1344.0],
+    "0.6": [768.0, 1280.0],
+    "0.68": [832.0, 1216.0],
+    "0.72": [832.0, 1152.0],
+    "0.78": [896.0, 1152.0],
+    "0.82": [896.0, 1088.0],
+    "0.88": [960.0, 1088.0],
+    "0.94": [960.0, 1024.0],
+    "1.0": [1024.0, 1024.0],
+    "1.07": [1024.0, 960.0],
+    "1.13": [1088.0, 960.0],
+    "1.21": [1088.0, 896.0],
+    "1.29": [1152.0, 896.0],
+    "1.38": [1152.0, 832.0],
+    "1.46": [1216.0, 832.0],
+    "1.67": [1280.0, 768.0],
+    "1.75": [1344.0, 768.0],
+    "2.0": [1408.0, 704.0],
+    "2.09": [1472.0, 704.0],
+    "2.4": [1536.0, 640.0],
+    "2.5": [1600.0, 640.0],
+    "3.0": [1728.0, 576.0],
+    "4.0": [2048.0, 512.0],
+}
+
+ASPECT_RATIO_512_BIN = {
+    "0.25": [256.0, 1024.0],
+    "0.28": [256.0, 928.0],
+    "0.32": [288.0, 896.0],
+    "0.33": [288.0, 864.0],
+    "0.35": [288.0, 832.0],
+    "0.4": [320.0, 800.0],
+    "0.42": [320.0, 768.0],
+    "0.48": [352.0, 736.0],
+    "0.5": [352.0, 704.0],
+    "0.52": [352.0, 672.0],
+    "0.57": [384.0, 672.0],
+    "0.6": [384.0, 640.0],
+    "0.68": [416.0, 608.0],
+    "0.72": [416.0, 576.0],
+    "0.78": [448.0, 576.0],
+    "0.82": [448.0, 544.0],
+    "0.88": [480.0, 544.0],
+    "0.94": [480.0, 512.0],
+    "1.0": [512.0, 512.0],
+    "1.07": [512.0, 480.0],
+    "1.13": [544.0, 480.0],
+    "1.21": [544.0, 448.0],
+    "1.29": [576.0, 448.0],
+    "1.38": [576.0, 416.0],
+    "1.46": [608.0, 416.0],
+    "1.67": [640.0, 384.0],
+    "1.75": [672.0, 384.0],
+    "2.0": [704.0, 352.0],
+    "2.09": [736.0, 352.0],
+    "2.4": [768.0, 320.0],
+    "2.5": [800.0, 320.0],
+    "3.0": [864.0, 288.0],
+    "4.0": [1024.0, 256.0],
+}

models/RewardStableDiffusion.py

@@ -0,0 +1,277 @@
+import inspect
+from typing import Callable, List, Optional, Union
+
+import torch
+from diffusers import StableDiffusionPipeline
+
+
+def freeze_params(params):
+    for param in params:
+        param.requires_grad = False
+
+
+class RewardStableDiffusion(StableDiffusionPipeline):
+    def __init__(
+        self,
+        vae,
+        text_encoder,
+        tokenizer,
+        unet,
+        scheduler,
+        safety_checker,
+        feature_extractor,
+        image_encoder=None,
+        requires_safety_checker: bool = True,
+        memsave=False,
+    ):
+        super().__init__(
+            vae,
+            text_encoder,
+            tokenizer,
+            unet,
+            scheduler,
+            safety_checker,
+            feature_extractor,
+            image_encoder,
+        )
+        # optionally enable memsave_torch
+        if memsave:
+            import memsave_torch.nn
+
+            self.vae = memsave_torch.nn.convert_to_memory_saving(self.vae)
+            self.unet = memsave_torch.nn.convert_to_memory_saving(self.unet)
+            self.text_encoder = memsave_torch.nn.convert_to_memory_saving(
+                self.text_encoder
+            )
+        # enable checkpointing
+        self.text_encoder.gradient_checkpointing_enable()
+        self.unet.enable_gradient_checkpointing()
+        self.vae.eval()
+        self.text_encoder.eval()
+        self.unet.eval()
+
+        # freeze diffusion parameters
+        freeze_params(self.vae.parameters())
+        freeze_params(self.unet.parameters())
+        freeze_params(self.text_encoder.parameters())
+
+    def decode_latents_tensors(self, latents):
+        latents = 1 / 0.18215 * latents
+        image = self.vae.decode(latents).sample
+        image = (image / 2 + 0.5).clamp(0, 1)
+        return image
+
+    def apply(
+        self,
+        latents: torch.Tensor,
+        prompt: Union[str, List[str]] = None,
+        text_embeddings=None,
+        image=None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        timesteps: Optional[List[int]] = None,
+        num_inference_steps: int = 1,
+        guidance_scale: float = 1.0,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+        callback_steps: Optional[int] = 1,
+    ) -> torch.Tensor:
+        # 0. Default height and width to unet
+        height = height or self.unet.config.sample_size * self.vae_scale_factor
+        width = width or self.unet.config.sample_size * self.vae_scale_factor
+        # to deal with lora scaling and other possible forward hooks
+
+        prompt_embeds = None
+        negative_prompt_embeds = None
+        ip_adapter_image = None
+        ip_adapter_image_embeds = None
+        callback_on_step_end_tensor_inputs = None
+        guidance_rescale = 0.0
+        clip_skip = None
+        cross_attention_kwargs = None
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            height,
+            width,
+            callback_steps,
+            negative_prompt,
+            prompt_embeds,
+            negative_prompt_embeds,
+            ip_adapter_image,
+            ip_adapter_image_embeds,
+            callback_on_step_end_tensor_inputs,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._guidance_rescale = guidance_rescale
+        self._clip_skip = clip_skip
+        self._cross_attention_kwargs = cross_attention_kwargs
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+
+        # 3. Encode input prompt
+        lora_scale = (
+            self.cross_attention_kwargs.get("scale", None)
+            if self.cross_attention_kwargs is not None
+            else None
+        )
+
+        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+            prompt,
+            device,
+            num_images_per_prompt,
+            self.do_classifier_free_guidance,
+            negative_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            lora_scale=lora_scale,
+            clip_skip=self.clip_skip,
+        )
+
+        # For classifier free guidance, we need to do two forward passes.
+        # Here we concatenate the unconditional and text embeddings into a single batch
+        # to avoid doing two forward passes
+        if self.do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
+
+        if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
+            image_embeds = self.prepare_ip_adapter_image_embeds(
+                ip_adapter_image,
+                ip_adapter_image_embeds,
+                device,
+                batch_size * num_images_per_prompt,
+                self.do_classifier_free_guidance,
+            )
+
+        # 4. Prepare timesteps
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler, num_inference_steps, device, timesteps
+        )
+
+        # 5. Prepare latent variables
+        num_channels_latents = self.unet.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 6.1 Add image embeds for IP-Adapter
+        added_cond_kwargs = (
+            {"image_embeds": image_embeds}
+            if (ip_adapter_image is not None or ip_adapter_image_embeds is not None)
+            else None
+        )
+
+        # 6.2 Optionally get Guidance Scale Embedding
+        timestep_cond = None
+        if self.unet.config.time_cond_proj_dim is not None:
+            guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(
+                batch_size * num_images_per_prompt
+            )
+            timestep_cond = self.get_guidance_scale_embedding(
+                guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
+            ).to(device=device, dtype=latents.dtype)
+
+        # 7. Denoising loop
+        self._num_timesteps = len(timesteps)
+        for i, t in enumerate(timesteps):
+            # expand the latents if we are doing classifier free guidance
+            latent_model_input = (
+                torch.cat([latents] * 2)
+                if self.do_classifier_free_guidance
+                else latents
+            )
+            latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+            # predict the noise residual
+            noise_pred = self.unet(
+                latent_model_input,
+                t,
+                encoder_hidden_states=prompt_embeds,
+                timestep_cond=timestep_cond,
+                added_cond_kwargs=added_cond_kwargs,
+                return_dict=False,
+            )[0]
+
+            # perform guidance
+            if self.do_classifier_free_guidance:
+                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                noise_pred = noise_pred_uncond + guidance_scale * (
+                    noise_pred_text - noise_pred_uncond
+                )
+
+            # compute the previous noisy sample x_t -> x_t-1
+            latents = self.scheduler.step(
+                noise_pred, t, latents, **extra_step_kwargs, return_dict=False
+            )[0]
+
+        image = self.decode_latents_tensors(latents)
+        return image
+
+
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    **kwargs,
+):
+    """
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+                Custom timesteps used to support arbitrary spacing between timesteps. If `None`, then the default
+                timestep spacing strategy of the scheduler is used. If `timesteps` is passed, `num_inference_steps`
+                must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(
+            inspect.signature(scheduler.set_timesteps).parameters.keys()
+        )
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps

models/RewardStableDiffusionXL.py

@@ -0,0 +1,320 @@
+from typing import List, Optional, Union
+
+import torch
+from diffusers import (AutoencoderKL, StableDiffusionXLPipeline,
+                       UNet2DConditionModel)
+from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import \
+    retrieve_timesteps
+from diffusers.schedulers import KarrasDiffusionSchedulers
+from transformers import (CLIPImageProcessor, CLIPTextModel,
+                          CLIPTextModelWithProjection, CLIPTokenizer,
+                          CLIPVisionModelWithProjection)
+
+
+def freeze_params(params):
+    for param in params:
+        param.requires_grad = False
+
+
+class RewardStableDiffusionXL(StableDiffusionXLPipeline):
+    def __init__(
+        self,
+        vae: AutoencoderKL,
+        text_encoder: CLIPTextModel,
+        text_encoder_2: CLIPTextModelWithProjection,
+        tokenizer: CLIPTokenizer,
+        tokenizer_2: CLIPTokenizer,
+        unet: UNet2DConditionModel,
+        scheduler: KarrasDiffusionSchedulers,
+        image_encoder: CLIPVisionModelWithProjection = None,
+        feature_extractor: CLIPImageProcessor = None,
+        force_zeros_for_empty_prompt: bool = True,
+        add_watermarker: bool = False,
+        is_hyper: bool = False,
+        memsave: bool = False,
+    ):
+        super().__init__(
+            vae,
+            text_encoder,
+            text_encoder_2,
+            tokenizer,
+            tokenizer_2,
+            unet,
+            scheduler,
+            image_encoder,
+            feature_extractor,
+            force_zeros_for_empty_prompt,
+            add_watermarker,
+        )
+        # optionally enable memsave_torch
+        if memsave:
+            import memsave_torch.nn
+
+            self.vae = memsave_torch.nn.convert_to_memory_saving(self.vae)
+            self.unet = memsave_torch.nn.convert_to_memory_saving(self.unet)
+            self.text_encoder = memsave_torch.nn.convert_to_memory_saving(
+                self.text_encoder
+            )
+            self.text_encoder_2 = memsave_torch.nn.convert_to_memory_saving(
+                self.text_encoder_2
+            )
+        # enable checkpointing
+        self.unet.enable_gradient_checkpointing()
+        self.vae.enable_gradient_checkpointing()
+        self.text_encoder.eval()
+        self.text_encoder_2.eval()
+        self.unet.eval()
+        self.vae.eval()
+        self.is_hyper = is_hyper
+
+        # freeze diffusion parameters
+        freeze_params(self.vae.parameters())
+        freeze_params(self.unet.parameters())
+        freeze_params(self.text_encoder.parameters())
+        freeze_params(self.text_encoder_2.parameters())
+
+    def decode_latents_tensors(self, latents):
+        latents = latents / self.vae.config.scaling_factor
+        image = self.vae.decode(latents).sample
+        image = (image / 2 + 0.5).clamp(0, 1)
+        return image
+
+    def apply(
+        self,
+        latents: torch.Tensor,
+        prompt: Union[str, List[str]] = None,
+        prompt_2: Optional[Union[str, List[str]]] = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 1,
+        guidance_scale: float = 0.0,
+        timesteps: List[int] = None,
+        denoising_end: Optional[float] = None,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        negative_prompt_2: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+    ) -> torch.Tensor:
+        if self.is_hyper:
+            timesteps = [800]
+        # 0. Default height and width to unet
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+
+        original_size = (height, width)
+        target_size = (height, width)
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            prompt_2,
+            height,
+            width,
+            callback_steps=1,
+        )
+
+        # 2. Define call parameters
+
+        self._guidance_scale = guidance_scale
+        self._clip_skip = 0
+        self._cross_attention_kwargs = None
+        self._denoising_end = denoising_end
+        self._interrupt = False
+
+        # 2. Define call parameters
+        batch_size = 1
+        device = self._execution_device
+
+        # 3. Encode input prompt
+        lora_scale = (
+            self.cross_attention_kwargs.get("scale", None)
+            if self.cross_attention_kwargs is not None
+            else None
+        )
+        prompt_embeds = None
+        negative_prompt_embeds = None
+        pooled_prompt_embeds = None
+        negative_pooled_prompt_embeds = None
+        (
+            prompt_embeds,
+            negative_prompt_embeds,
+            pooled_prompt_embeds,
+            negative_pooled_prompt_embeds,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            do_classifier_free_guidance=self.do_classifier_free_guidance,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt_2,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            lora_scale=lora_scale,
+            clip_skip=self.clip_skip,
+        )
+
+        # 4. Prepare timesteps
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler, num_inference_steps, device, timesteps
+        )
+
+        num_channels_latents = self.unet.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 7. Prepare added time ids & embeddings
+        add_text_embeds = pooled_prompt_embeds
+        if self.text_encoder_2 is None:
+            text_encoder_projection_dim = int(pooled_prompt_embeds.shape[-1])
+        else:
+            text_encoder_projection_dim = self.text_encoder_2.config.projection_dim
+
+        add_time_ids = self._get_add_time_ids(
+            original_size,
+            (0, 0),
+            target_size,
+            dtype=prompt_embeds.dtype,
+            text_encoder_projection_dim=text_encoder_projection_dim,
+        )
+        negative_add_time_ids = add_time_ids
+
+        if self.do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0)
+            add_text_embeds = torch.cat(
+                [negative_pooled_prompt_embeds, add_text_embeds], dim=0
+            )
+            add_time_ids = torch.cat([negative_add_time_ids, add_time_ids], dim=0)
+
+        prompt_embeds = prompt_embeds.to(device)
+        add_text_embeds = add_text_embeds.to(device)
+        add_time_ids = add_time_ids.to(device).repeat(
+            batch_size * num_images_per_prompt, 1
+        )
+
+        # 8. Denoising loop
+        num_warmup_steps = max(
+            len(timesteps) - num_inference_steps * self.scheduler.order, 0
+        )
+
+        # 8.1 Apply denoising_end
+        if (
+            self.denoising_end is not None
+            and isinstance(self.denoising_end, float)
+            and self.denoising_end > 0
+            and self.denoising_end < 1
+        ):
+            discrete_timestep_cutoff = int(
+                round(
+                    self.scheduler.config.num_train_timesteps
+                    - (self.denoising_end * self.scheduler.config.num_train_timesteps)
+                )
+            )
+            num_inference_steps = len(
+                list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))
+            )
+            timesteps = timesteps[:num_inference_steps]
+
+        # 9. Optionally get Guidance Scale Embedding
+        timestep_cond = None
+        if self.unet.config.time_cond_proj_dim is not None:
+            guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(
+                batch_size * num_images_per_prompt
+            )
+            timestep_cond = self.get_guidance_scale_embedding(
+                guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
+            ).to(device=device, dtype=latents.dtype)
+
+        self._num_timesteps = len(timesteps)
+
+        # 8. Denoising loop
+        # 8.1 Apply denoising_end
+        if (
+            self.denoising_end is not None
+            and isinstance(self.denoising_end, float)
+            and self.denoising_end > 0
+            and self.denoising_end < 1
+        ):
+            discrete_timestep_cutoff = int(
+                round(
+                    self.scheduler.config.num_train_timesteps
+                    - (self.denoising_end * self.scheduler.config.num_train_timesteps)
+                )
+            )
+            num_inference_steps = len(
+                list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))
+            )
+            timesteps = timesteps[:num_inference_steps]
+
+        # 9. Optionally get Guidance Scale Embedding
+        timestep_cond = None
+
+        self._num_timesteps = len(timesteps)
+        for i, t in enumerate(timesteps):
+            if self._interrupt:
+                continue
+            # expand the latents if we are doing classifier free guidance
+            latent_model_input = (
+                torch.cat([latents] * 2)
+                if self.do_classifier_free_guidance
+                else latents
+            )
+
+            latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+            # predict the noise residual
+            added_cond_kwargs = {
+                "text_embeds": add_text_embeds,
+                "time_ids": add_time_ids,
+            }
+            noise_pred = self.unet(
+                latent_model_input,
+                t,
+                encoder_hidden_states=prompt_embeds,
+                timestep_cond=timestep_cond,
+                cross_attention_kwargs=self.cross_attention_kwargs,
+                added_cond_kwargs=added_cond_kwargs,
+                return_dict=False,
+            )[0]
+
+            # perform guidance
+            if self.do_classifier_free_guidance:
+                noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                noise_pred = noise_pred_uncond + self.guidance_scale * (
+                    noise_pred_text - noise_pred_uncond
+                )
+
+            # compute the previous noisy sample x_t -> x_t-1
+            latents = self.scheduler.step(
+                noise_pred, t, latents, **extra_step_kwargs, return_dict=False
+            )[0]
+
+        if self.is_hyper:
+            latents = latents.to(torch.float32)
+            image = self.decode_latents_tensors(latents)
+            image = image.to(torch.float16)
+        else:
+            image = self.decode_latents_tensors(latents)
+
+        # apply watermark if available
+        if self.watermark is not None:
+            image = self.watermark.apply_watermark(image)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        return image

models/__init__.py

@@ -0,0 +1 @@
+from .utils import get_model

models/utils.py

@@ -0,0 +1,109 @@
+import logging
+
+import torch
+from diffusers import (AutoencoderKL, DDPMScheduler,
+                       EulerAncestralDiscreteScheduler, LCMScheduler,
+                       Transformer2DModel, UNet2DConditionModel)
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
+
+from models.RewardPixart import RewardPixartPipeline, freeze_params
+from models.RewardStableDiffusion import RewardStableDiffusion
+from models.RewardStableDiffusionXL import RewardStableDiffusionXL
+
+
+def get_model(
+    model_name: str,
+    dtype: torch.dtype,
+    device: torch.device,
+    cache_dir: str,
+    memsave: bool = False,
+):
+    logging.info(f"Loading model: {model_name}")
+    if model_name == "sd-turbo":
+        pipe = RewardStableDiffusion.from_pretrained(
+            "stabilityai/sd-turbo",
+            torch_dtype=dtype,
+            variant="fp16",
+            cache_dir=cache_dir,
+            memsave=memsave,
+        )
+        pipe = pipe.to(device, dtype)
+    elif model_name == "sdxl-turbo":
+        vae = AutoencoderKL.from_pretrained(
+            "madebyollin/sdxl-vae-fp16-fix",
+            torch_dtype=torch.float16,
+            cache_dir=cache_dir,
+        )
+        pipe = RewardStableDiffusionXL.from_pretrained(
+            "stabilityai/sdxl-turbo",
+            vae=vae,
+            torch_dtype=dtype,
+            variant="fp16",
+            use_safetensors=True,
+            cache_dir=cache_dir,
+            memsave=memsave,
+        )
+        pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(
+            pipe.scheduler.config, timestep_spacing="trailing"
+        )
+        pipe = pipe.to(device, dtype)
+    elif model_name == "pixart":
+        pipe = RewardPixartPipeline.from_pretrained(
+            "PixArt-alpha/PixArt-XL-2-1024-MS",
+            torch_dtype=dtype,
+            cache_dir=cache_dir,
+            memsave=memsave,
+        )
+        pipe.transformer = Transformer2DModel.from_pretrained(
+            "PixArt-alpha/PixArt-Alpha-DMD-XL-2-512x512",
+            subfolder="transformer",
+            torch_dtype=dtype,
+            cache_dir=cache_dir,
+        )
+        pipe.scheduler = DDPMScheduler.from_pretrained(
+            "PixArt-alpha/PixArt-Alpha-DMD-XL-2-512x512",
+            subfolder="scheduler",
+            cache_dir=cache_dir,
+        )
+
+        # speed-up T5
+        pipe.text_encoder.to_bettertransformer()
+        pipe.transformer.eval()
+        freeze_params(pipe.transformer.parameters())
+        pipe.transformer.enable_gradient_checkpointing()
+        pipe = pipe.to(device)
+    elif model_name == "hyper-sd":
+        base_model_id = "stabilityai/stable-diffusion-xl-base-1.0"
+        repo_name = "ByteDance/Hyper-SD"
+        ckpt_name = "Hyper-SDXL-1step-Unet.safetensors"
+        # Load model.
+        unet = UNet2DConditionModel.from_config(
+            base_model_id, subfolder="unet", cache_dir=cache_dir
+        ).to(device, dtype)
+        unet.load_state_dict(
+            load_file(
+                hf_hub_download(repo_name, ckpt_name, cache_dir=cache_dir),
+                device="cuda",
+            )
+        )
+        pipe = RewardStableDiffusionXL.from_pretrained(
+            base_model_id,
+            unet=unet,
+            torch_dtype=dtype,
+            variant="fp16",
+            cache_dir=cache_dir,
+            is_hyper=True,
+            memsave=memsave,
+        )
+        # Use LCM scheduler instead of ddim scheduler to support specific timestep number inputs
+        pipe.scheduler = LCMScheduler.from_config(
+            pipe.scheduler.config, cache_dir=cache_dir
+        )
+        pipe = pipe.to(device, dtype)
+        # upcast vae
+        pipe.vae = pipe.vae.to(dtype=torch.float32)
+        # pipe.enable_sequential_cpu_offload()
+    else:
+        raise ValueError(f"Unknown model name: {model_name}")
+    return pipe

rewards/__init__.py

@@ -0,0 +1 @@
+from .utils import clip_img_transform, get_reward_losses

rewards/aesthetic.py

@@ -0,0 +1,118 @@
+import os
+
+import clip
+import pytorch_lightning as pl
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from rewards.base_reward import BaseRewardLoss
+
+
+class AestheticLoss(BaseRewardLoss):
+    """CLIP reward loss function for optimization."""
+
+    def __init__(
+        self,
+        weigthing: float,
+        dtype: torch.dtype,
+        device: torch.device,
+        cache_dir: str,
+        memsave: bool = False,
+    ):
+        self.clip_model, self.preprocess_fn = clip.load(
+            "ViT-L/14", device=device, download_root=cache_dir
+        )
+        self.clip_model = self.clip_model.to(device, dtype=dtype)
+        self.mlp = MLP(768).to(device, dtype=dtype)
+        s = torch.load(
+            f"{os.getcwd()}/ckpts/aesthetic-model.pth"
+        )  # load the model you trained previously or the model available in this repo
+        self.mlp.load_state_dict(s)
+        self.clip_model.eval()
+        if memsave:
+            import memsave_torch.nn
+
+            self.mlp = memsave_torch.nn.convert_to_memory_saving(self.mlp)
+            self.clip_model = memsave_torch.nn.convert_to_memory_saving(
+                self.clip_model
+            ).to(device, dtype=dtype)
+
+        self.freeze_parameters(self.clip_model.parameters())
+        self.freeze_parameters(self.mlp.parameters())
+        super().__init__("Aesthetic", weigthing)
+
+    def get_image_features(self, image: torch.Tensor) -> torch.Tensor:
+        with torch.autocast("cuda"):
+            clip_img_features = self.clip_model.encode_image(image)
+            l2 = torch.norm(clip_img_features, p=2, dim=-1, keepdim=True)
+            l2 = torch.where(
+                l2 == 0,
+                torch.tensor(
+                    1.0, device=clip_img_features.device, dtype=clip_img_features.dtype
+                ),
+                l2,
+            )
+            clip_img_features = clip_img_features / l2
+        return clip_img_features
+
+    def get_text_features(self, prompt: str) -> torch.Tensor:
+        return None
+
+    def compute_loss(
+        self, image_features: torch.Tensor, text_features: torch.Tensor
+    ) -> torch.Tensor:
+        return None
+
+    def __call__(self, image: torch.Tensor, prompt: torch.Tensor) -> torch.Tensor:
+        if self.memsave:
+            image = image.to(torch.float32)
+        image_features = self.get_image_features(image)
+
+        image_features_normed = self.process_features(image_features.to(torch.float16))
+
+        aesthetic_loss = 10.0 - self.mlp(image_features_normed).mean()
+        return aesthetic_loss
+
+
+class MLP(pl.LightningModule):
+    def __init__(self, input_size, xcol="emb", ycol="avg_rating"):
+        super().__init__()
+        self.input_size = input_size
+        self.xcol = xcol
+        self.ycol = ycol
+        self.layers = nn.Sequential(
+            nn.Linear(self.input_size, 1024),
+            # nn.ReLU(),
+            nn.Dropout(0.2),
+            nn.Linear(1024, 128),
+            # nn.ReLU(),
+            nn.Dropout(0.2),
+            nn.Linear(128, 64),
+            # nn.ReLU(),
+            nn.Dropout(0.1),
+            nn.Linear(64, 16),
+            # nn.ReLU(),
+            nn.Linear(16, 1),
+        )
+
+    def forward(self, x):
+        return self.layers(x)
+
+    def training_step(self, batch, batch_idx):
+        x = batch[self.xcol]
+        y = batch[self.ycol].reshape(-1, 1)
+        x_hat = self.layers(x)
+        loss = F.mse_loss(x_hat, y)
+        return loss
+
+    def validation_step(self, batch, batch_idx):
+        x = batch[self.xcol]
+        y = batch[self.ycol].reshape(-1, 1)
+        x_hat = self.layers(x)
+        loss = F.mse_loss(x_hat, y)
+        return loss
+
+    def configure_optimizers(self):
+        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
+        return optimizer

rewards/base_reward.py

@@ -0,0 +1,46 @@
+from abc import ABC, abstractmethod
+
+import torch
+
+
+class BaseRewardLoss(ABC):
+    """
+    Base class for reward functions implementing a differentiable reward function for optimization.
+    """
+
+    def __init__(self, name: str, weighting: float):
+        self.name = name
+        self.weighting = weighting
+
+    @staticmethod
+    def freeze_parameters(params: torch.nn.ParameterList):
+        for param in params:
+            param.requires_grad = False
+
+    @abstractmethod
+    def get_image_features(self, image: torch.Tensor) -> torch.Tensor:
+        pass
+
+    @abstractmethod
+    def get_text_features(self, prompt: str) -> torch.Tensor:
+        pass
+
+    @abstractmethod
+    def compute_loss(
+        self, image_features: torch.Tensor, text_features: torch.Tensor
+    ) -> torch.Tensor:
+        pass
+
+    def process_features(self, features: torch.Tensor) -> torch.Tensor:
+        features_normed = features / features.norm(dim=-1, keepdim=True)
+        return features_normed
+
+    def __call__(self, image: torch.Tensor, prompt: str) -> torch.Tensor:
+        image_features = self.get_image_features(image)
+        text_features = self.get_text_features(prompt)
+
+        image_features_normed = self.process_features(image_features)
+        text_features_normed = self.process_features(text_features)
+
+        loss = self.compute_loss(image_features_normed, text_features_normed)
+        return loss

rewards/clip.py

@@ -0,0 +1,54 @@
+import torch
+from transformers import CLIPModel
+
+from rewards.base_reward import BaseRewardLoss
+
+
+class CLIPLoss(BaseRewardLoss):
+    """CLIP reward loss function for optimization."""
+
+    def __init__(
+        self,
+        weigthing: float,
+        dtype: torch.dtype,
+        device: torch.device,
+        cache_dir: str,
+        tokenizer,
+        memsave: bool = False,
+    ):
+        self.tokenizer = tokenizer
+        self.clip_model = CLIPModel.from_pretrained(
+            "laion/CLIP-ViT-H-14-laion2B-s32B-b79K",
+            cache_dir=cache_dir,
+        )
+        # freeze all models parameters
+        if memsave:
+            import memsave_torch.nn
+
+            self.clip_model = memsave_torch.nn.convert_to_memory_saving(self.clip_model)
+        self.clip_model = self.clip_model.to(device, dtype=dtype)
+        self.clip_model.eval()
+        self.freeze_parameters(self.clip_model.parameters())
+        super().__init__("CLIP", weigthing)
+        self.clip_model.gradient_checkpointing_enable()
+
+    def get_image_features(self, image: torch.Tensor) -> torch.Tensor:
+        clip_img_features = self.clip_model.get_image_features(image)
+        return clip_img_features
+
+    def get_text_features(self, prompt: str) -> torch.Tensor:
+        prompt_token = self.tokenizer(
+            prompt, return_tensors="pt", padding=True, max_length=77, truncation=True
+        ).to("cuda")
+        clip_text_features = self.clip_model.get_text_features(**prompt_token)
+        return clip_text_features
+
+    def compute_loss(
+        self, image_features: torch.Tensor, text_features: torch.Tensor
+    ) -> torch.Tensor:
+        clip_loss = (
+            100
+            - (image_features @ text_features.T).mean()
+            * self.clip_model.logit_scale.exp()
+        )
+        return clip_loss

rewards/hps.py

@@ -0,0 +1,57 @@
+import huggingface_hub
+import torch
+from hpsv2.src.open_clip import create_model, get_tokenizer
+
+from rewards.base_reward import BaseRewardLoss
+
+
+class HPSLoss(BaseRewardLoss):
+    """HPS reward loss function for optimization."""
+
+    def __init__(
+        self,
+        weighting: float,
+        dtype: torch.dtype,
+        device: torch.device,
+        cache_dir: str,
+        memsave: bool = False,
+    ):
+        self.hps_model = create_model(
+            "ViT-H-14",
+            "laion2B-s32B-b79K",
+            precision=dtype,
+            device=device,
+            cache_dir=cache_dir,
+        )
+        checkpoint_path = huggingface_hub.hf_hub_download(
+            "xswu/HPSv2", "HPS_v2.1_compressed.pt", cache_dir=cache_dir
+        )
+        self.hps_model.load_state_dict(
+            torch.load(checkpoint_path, map_location=device)["state_dict"]
+        )
+        self.hps_tokenizer = get_tokenizer("ViT-H-14")
+        if memsave:
+            import memsave_torch.nn
+
+            self.hps_model = memsave_torch.nn.convert_to_memory_saving(self.hps_model)
+        self.hps_model = self.hps_model.to(device, dtype=dtype)
+        self.hps_model.eval()
+        self.freeze_parameters(self.hps_model.parameters())
+        super().__init__("HPS", weighting)
+        self.hps_model.set_grad_checkpointing(True)
+
+    def get_image_features(self, image: torch.Tensor) -> torch.Tensor:
+        hps_image_features = self.hps_model.encode_image(image)
+        return hps_image_features
+
+    def get_text_features(self, prompt: str) -> torch.Tensor:
+        hps_text = self.hps_tokenizer(prompt).to("cuda")
+        hps_text_features = self.hps_model.encode_text(hps_text)
+        return hps_text_features
+
+    def compute_loss(
+        self, image_features: torch.Tensor, text_features: torch.Tensor
+    ) -> torch.Tensor:
+        logits_per_image = image_features @ text_features.T
+        hps_loss = 1 - torch.diagonal(logits_per_image)[0]
+        return hps_loss

rewards/imagereward.py

@@ -0,0 +1,61 @@
+import ImageReward as RM
+import torch
+
+from rewards.base_reward import BaseRewardLoss
+
+
+class ImageRewardLoss:
+    """Image reward loss for optimization."""
+
+    def __init__(
+        self,
+        weighting: float,
+        dtype: torch.dtype,
+        device: torch.device,
+        cache_dir: str,
+        memsave: bool = False,
+    ):
+        self.name = "ImageReward"
+        self.weighting = weighting
+        self.dtype = dtype
+        self.imagereward_model = RM.load("ImageReward-v1.0", download_root=cache_dir)
+        self.imagereward_model = self.imagereward_model.to(
+            device=device, dtype=self.dtype
+        )
+        self.imagereward_model.eval()
+        BaseRewardLoss.freeze_parameters(self.imagereward_model.parameters())
+
+    def __call__(self, image: torch.Tensor, prompt: str) -> torch.Tensor:
+        imagereward_score = self.score_diff(prompt, image)
+        return (2 - imagereward_score).mean()
+
+    def score_diff(self, prompt, image):
+        # text encode
+        text_input = self.imagereward_model.blip.tokenizer(
+            prompt,
+            padding="max_length",
+            truncation=True,
+            max_length=35,
+            return_tensors="pt",
+        ).to(self.imagereward_model.device)
+        image_embeds = self.imagereward_model.blip.visual_encoder(image)
+
+        # text encode cross attention with image
+        image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(
+            self.imagereward_model.device
+        )
+        text_output = self.imagereward_model.blip.text_encoder(
+            text_input.input_ids,
+            attention_mask=text_input.attention_mask,
+            encoder_hidden_states=image_embeds,
+            encoder_attention_mask=image_atts,
+            return_dict=True,
+        )
+
+        txt_features = text_output.last_hidden_state[:, 0, :].to(
+            self.imagereward_model.device, dtype=self.dtype
+        )
+        rewards = self.imagereward_model.mlp(txt_features)
+        rewards = (rewards - self.imagereward_model.mean) / self.imagereward_model.std
+
+        return rewards

rewards/pickscore.py

@@ -0,0 +1,55 @@
+import torch
+from transformers import AutoModel
+
+from rewards.base_reward import BaseRewardLoss
+
+
+class PickScoreLoss(BaseRewardLoss):
+    """PickScore reward loss function for optimization."""
+
+    def __init__(
+        self,
+        weighting: float,
+        dtype: torch.dtype,
+        device: torch.device,
+        cache_dir: str,
+        tokenizer,
+        memsave: bool = False,
+    ):
+        self.tokenizer = tokenizer
+        self.pickscore_model = AutoModel.from_pretrained(
+            "yuvalkirstain/PickScore_v1", cache_dir=cache_dir
+        ).eval()
+        if memsave:
+            import memsave_torch.nn
+
+            self.pickscore_model = memsave_torch.nn.convert_to_memory_saving(
+                self.pickscore_model
+            )
+        self.pickscore_model = self.pickscore_model.to(device, dtype=dtype)
+        self.freeze_parameters(self.pickscore_model.parameters())
+        super().__init__("PickScore", weighting)
+        self.pickscore_model._set_gradient_checkpointing(True)
+
+    def get_image_features(self, image) -> torch.Tensor:
+        reward_img_features = self.pickscore_model.get_image_features(image)
+        return reward_img_features
+
+    def get_text_features(self, prompt: str) -> torch.Tensor:
+        prompt_token = self.tokenizer(
+            prompt, return_tensors="pt", padding=True, max_length=77, truncation=True
+        ).to("cuda")
+        reward_text_features = self.pickscore_model.get_text_features(**prompt_token)
+        return reward_text_features
+
+    def compute_loss(
+        self, image_features: torch.Tensor, text_features: torch.Tensor
+    ) -> torch.Tensor:
+        pickscore_loss = (
+            30
+            - (
+                self.pickscore_model.logit_scale.exp()
+                * (image_features @ text_features.T)
+            ).mean()
+        )
+        return pickscore_loss

rewards/utils.py

@@ -0,0 +1,79 @@
+from typing import Any, List
+
+import torch
+from torchvision.transforms import (CenterCrop, Compose, InterpolationMode,
+                                    Normalize, Resize)
+from transformers import AutoProcessor
+
+from rewards.aesthetic import AestheticLoss
+from rewards.base_reward import BaseRewardLoss
+from rewards.clip import CLIPLoss
+from rewards.hps import HPSLoss
+from rewards.imagereward import ImageRewardLoss
+from rewards.pickscore import PickScoreLoss
+
+
+def get_reward_losses(
+    args: Any, dtype: torch.dtype, device: torch.device, cache_dir: str
+) -> List[BaseRewardLoss]:
+    if args.enable_clip or args.enable_pickscore:
+        tokenizer = AutoProcessor.from_pretrained(
+            "laion/CLIP-ViT-H-14-laion2B-s32B-b79K", cache_dir=cache_dir
+        )
+    reward_losses = []
+    if args.enable_hps:
+        reward_losses.append(
+            HPSLoss(args.hps_weighting, dtype, device, cache_dir, memsave=args.memsave)
+        )
+    if args.enable_imagereward:
+        reward_losses.append(
+            ImageRewardLoss(
+                args.imagereward_weighting,
+                dtype,
+                device,
+                cache_dir,
+                memsave=args.memsave,
+            )
+        )
+    if args.enable_clip:
+        reward_losses.append(
+            CLIPLoss(
+                args.clip_weighting,
+                dtype,
+                device,
+                cache_dir,
+                tokenizer,
+                memsave=args.memsave,
+            )
+        )
+    if args.enable_pickscore:
+        reward_losses.append(
+            PickScoreLoss(
+                args.pickscore_weighting,
+                dtype,
+                device,
+                cache_dir,
+                tokenizer,
+                memsave=args.memsave,
+            )
+        )
+    if args.enable_aesthetic:
+        reward_losses.append(
+            AestheticLoss(
+                args.aesthetic_weighting, dtype, device, cache_dir, memsave=args.memsave
+            )
+        )
+    return reward_losses
+
+
+def clip_img_transform(size: int = 224):
+    return Compose(
+        [
+            Resize(size, interpolation=InterpolationMode.BICUBIC),
+            CenterCrop(size),
+            Normalize(
+                (0.48145466, 0.4578275, 0.40821073),
+                (0.26862954, 0.26130258, 0.27577711),
+            ),
+        ]
+    )

training/__init__.py

@@ -0,0 +1,2 @@
+from .optim import get_optimizer
+from .trainer import LatentNoiseTrainer

training/optim.py

@@ -0,0 +1,21 @@
+import torch
+
+
+def get_optimizer(
+    optimizer_name: str, latents: torch.Tensor, lr: float, nesterov: bool
+):
+    if optimizer_name == "adam":
+        optimizer = torch.optim.Adam([latents], lr=lr, eps=1e-2)
+    elif optimizer_name == "sgd":
+        optimizer = torch.optim.SGD([latents], lr=lr, nesterov=nesterov, momentum=0.9)
+    elif optimizer_name == "lbfgs":
+        optimizer = torch.optim.LBFGS(
+            [latents],
+            lr=lr,
+            max_iter=10,
+            history_size=3,
+            line_search_fn="strong_wolfe",
+        )
+    else:
+        raise ValueError(f"Unknown optimizer {optimizer_name}")
+    return optimizer

training/trainer.py

@@ -0,0 +1,125 @@
+import logging
+import math
+from typing import Dict, List, Optional, Tuple
+
+import PIL
+import PIL.Image
+import torch
+from diffusers import DiffusionPipeline
+
+from rewards import clip_img_transform
+from rewards.base_reward import BaseRewardLoss
+
+
+class LatentNoiseTrainer:
+    """Trainer for optimizing latents with reward losses."""
+
+    def __init__(
+        self,
+        reward_losses: List[BaseRewardLoss],
+        model: DiffusionPipeline,
+        n_iters: int,
+        n_inference_steps: int,
+        seed: int,
+        no_optim: bool = False,
+        regularize: bool = True,
+        regularization_weight: float = 0.01,
+        grad_clip: float = 0.1,
+        log_metrics: bool = True,
+        save_all_images: bool = False,
+        imageselect: bool = False,
+        device: torch.device = torch.device("cuda"),
+    ):
+        self.reward_losses = reward_losses
+        self.model = model
+        self.n_iters = n_iters
+        self.n_inference_steps = n_inference_steps
+        self.seed = seed
+        self.no_optim = no_optim
+        self.regularize = regularize
+        self.regularization_weight = regularization_weight
+        self.grad_clip = grad_clip
+        self.log_metrics = log_metrics
+        self.save_all_images = save_all_images
+        self.imageselect = imageselect
+        self.device = device
+        self.preprocess_fn = clip_img_transform(224)
+
+    def train(
+        self,
+        latents: torch.Tensor,
+        prompt: str,
+        optimizer: torch.optim.Optimizer,
+        save_dir: Optional[str] = None,
+    ) -> Tuple[PIL.Image.Image, Dict[str, float], Dict[str, float]]:
+        logging.info(f"Optimizing latents for prompt '{prompt}'.")
+        best_loss = torch.inf
+        best_image = None
+        initial_rewards = None
+        best_rewards = None
+        latent_dim = math.prod(latents.shape[1:])
+        for iteration in range(self.n_iters):
+            to_log = ""
+            rewards = {}
+            optimizer.zero_grad()
+            generator = torch.Generator("cuda").manual_seed(self.seed)
+            if self.imageselect:
+                new_latents = torch.randn_like(
+                    latents, device=self.device, dtype=latents.dtype
+                )
+                image = self.model.apply(
+                    new_latents,
+                    prompt,
+                    generator=generator,
+                    num_inference_steps=self.n_inference_steps,
+                )
+            else:
+                image = self.model.apply(
+                    latents,
+                    prompt,
+                    generator=generator,
+                    num_inference_steps=self.n_inference_steps,
+                )
+            if self.no_optim:
+                best_image = image
+                break
+
+            total_loss = 0
+            preprocessed_image = self.preprocess_fn(image)
+            for reward_loss in self.reward_losses:
+                loss = reward_loss(preprocessed_image, prompt)
+                to_log += f"{reward_loss.name}: {loss.item():.4f}, "
+                total_loss += loss * reward_loss.weighting
+                rewards[reward_loss.name] = loss.item()
+            rewards["total"] = total_loss.item()
+            to_log += f"Total: {total_loss.item():.4f}"
+            total_reward_loss = total_loss.item()
+            if self.regularize:
+                # compute in fp32 to avoid overflow
+                latent_norm = torch.linalg.vector_norm(latents).to(torch.float32)
+                log_norm = torch.log(latent_norm)
+                regularization = self.regularization_weight * (
+                    0.5 * latent_norm**2 - (latent_dim - 1) * log_norm
+                )
+                to_log += f", Latent norm: {latent_norm.item()}"
+                rewards["norm"] = latent_norm.item()
+                total_loss += regularization.to(total_loss.dtype)
+            if self.log_metrics:
+                logging.info(f"Iteration {iteration}: {to_log}")
+            if initial_rewards is None:
+                initial_rewards = rewards
+            if total_reward_loss < best_loss:
+                best_loss = total_reward_loss
+                best_image = image
+                best_rewards = rewards
+            if iteration != self.n_iters - 1 and not self.imageselect:
+                total_loss.backward()
+                torch.nn.utils.clip_grad_norm_(latents, self.grad_clip)
+                optimizer.step()
+            if self.save_all_images:
+                image_numpy = image.detach().cpu().permute(0, 2, 3, 1).float().numpy()
+                image_pil = DiffusionPipeline.numpy_to_pil(image_numpy)[0]
+                image_pil.save(f"{save_dir}/{iteration}.png")
+        image_numpy = best_image.detach().cpu().permute(0, 2, 3, 1).float().numpy()
+        image_pil = DiffusionPipeline.numpy_to_pil(image_numpy)[0]
+        return image_pil, initial_rewards, best_rewards