lrm/lrm_15/trainer/models/sd15_preference_model.py

from dataclasses import dataclass
import torch
from torch import nn
from torchvision import transforms
from PIL import Image
from diffusers import AutoencoderKL, DDPMScheduler
from transformers import CLIPTextModel, CLIPTokenizer, CLIPImageProcessor
import time
import os
from io import BytesIO
from trainer.models.base_model import BaseModelConfig
from trainer.models.unet_2d_condition_reward import UNet2DConditionModel

from accelerate.logging import get_logger
logger = get_logger(__name__)

@dataclass
class SD15PreferenceModelConfig(BaseModelConfig):
    _target_: str = "trainer.models.sd15_preference_model.SD15PreferenceModel"
    pretrained_model_name_or_path: str = 'runwayml/stable-diffusion-v1-5'
    clip_ckpt_path: str = 'openai/clip-vit-large-patch14/pytorch_model.bin'
    vae_path: str = ""
    vision_embed_dim: int = 1280
    text_embed_dim: int = 768
    projection_dim: int = 768
    logit_scale_init_value: float = 2.6592
    freeze_text_encoder: bool = False
    multi_scale: bool = True
    multi_scale_cfg: bool = False
    guidance_scale: float = 1.0


class SD15PreferenceModel(nn.Module):
    def __init__(self, cfg: SD15PreferenceModelConfig):
        super().__init__()
        # diffusion models
        self.tokenizer = CLIPTokenizer.from_pretrained(cfg.pretrained_model_name_or_path, subfolder="tokenizer")
        self.text_encoder = CLIPTextModel.from_pretrained(cfg.pretrained_model_name_or_path, subfolder="text_encoder")
        if cfg.vae_path != "":
            self.vae = AutoencoderKL.from_pretrained(cfg.vae_path)
        else:
            self.vae = AutoencoderKL.from_pretrained(cfg.pretrained_model_name_or_path, subfolder="vae")
        self.scheduler = DDPMScheduler.from_pretrained(cfg.pretrained_model_name_or_path, subfolder="scheduler")
        self.unet = UNet2DConditionModel.from_pretrained(cfg.pretrained_model_name_or_path, subfolder="unet")
        # self.pipeline = StableDiffusionPipeline.from_pretrained(pretrained_model_name_or_path)
        # self.image_processor = CLIPImageProcessor.from_pretrained(pretrained_model_name_or_path)
        
        # global pooling layer
        self.avg_pool = nn.AdaptiveAvgPool2d((1, 1))
        self.cfg = cfg
        
        # projection layers
        if cfg.multi_scale:
            self.visual_projection = nn.Linear(4800, cfg.projection_dim, bias=False)

        else:  
            self.visual_projection = nn.Linear(cfg.vision_embed_dim, cfg.projection_dim, bias=False)
        nn.init.normal_(self.visual_projection.weight, std=0.02)
        
        self.text_projection = nn.Linear(cfg.text_embed_dim, cfg.projection_dim, bias=False)
        # load text projections from openai/clip-vit-large-patch14
        clip_ckpt = torch.load(cfg.clip_ckpt_path)
        self.text_projection.weight.data = clip_ckpt['text_projection.weight'].contiguous()
        
        self.logit_scale = nn.Parameter(torch.ones([]) * cfg.logit_scale_init_value)
        
        self.vae.requires_grad_(False)
        if cfg.freeze_text_encoder:
            self.text_encoder.requires_grad_(False)
        
        self.val_transform = transforms.Compose(
            [
                transforms.Resize((512, 512), interpolation=transforms.InterpolationMode.BILINEAR),
                transforms.ToTensor(),
                transforms.Normalize([0.5], [0.5]),
            ]
        )
        self.do_classifier_free_guidance = self.cfg.guidance_scale > 1.0 or self.cfg.guidance_scale < 1.0
        if self.do_classifier_free_guidance:
            # generate negative prompt ids
            self.neg_prompt_ids = self.tokenizer(
                [""],
                return_tensors="pt",
                padding="max_length",
                truncation=True,
                max_length=self.tokenizer.model_max_length,
            ).input_ids
    
    def get_text_features(self, text_inputs=None):
        if self.do_classifier_free_guidance:
            text_inputs = torch.cat([text_inputs, self.neg_prompt_ids.repeat(text_inputs.shape[0], 1).to(text_inputs.device)], dim=0)
            
        outputs = self.text_encoder(text_inputs, return_dict=False)
        encoder_hidden_states = outputs[0]
        pooled_output = outputs[1]
        
        if self.do_classifier_free_guidance:
            pooled_output_text, pooled_output_ucond = pooled_output.chunk(2, dim=0)
            text_features = self.text_projection(pooled_output_text)
        else:
            text_features = self.text_projection(pooled_output)
        return encoder_hidden_states, text_features
    
    def get_image_features(self, encoder_hidden_states=None, image_inputs=None, time_cond=None, generator=None):
        with torch.no_grad():
            latents = self.vae.encode(image_inputs).latent_dist.sample()
            # latents = latents.to(dtype=self.unet.dtype)
            latents = latents * self.vae.config.scaling_factor
  
        if generator is not None:
            noise = torch.randn(latents.size(), generator=generator, dtype=latents.dtype, device=latents.device)
        else:
            noise = torch.randn_like(latents)

        noisy_latents = self.scheduler.add_noise(latents, noise, time_cond)
        
        if self.do_classifier_free_guidance:
            noisy_latents = torch.cat([noisy_latents] * 2, dim=0)
            time_cond = torch.cat([time_cond] * 2, dim=0)

        mid_output, down_block_res_samples = self.unet(noisy_latents, time_cond, encoder_hidden_states=encoder_hidden_states, return_dict=False, use_up_blocks=False)
        
        if self.cfg.multi_scale:
            first_stage_output = down_block_res_samples[2]  # [320, 64, 64]
            second_stage_output = down_block_res_samples[5]  # [640, 32, 32]
            third_stage_output = down_block_res_samples[8]  # [1280, 16, 16]
            fourth_stage_output = down_block_res_samples[11]  # [1280, 8, 8]

            pooled_first_stage_output = self.avg_pool(first_stage_output).squeeze(dim=[2,3])
            pooled_second_stage_output = self.avg_pool(second_stage_output).squeeze(dim=[2,3])
            pooled_third_stage_output = self.avg_pool(third_stage_output).squeeze(dim=[2,3])
            pooled_fourth_stage_output = self.avg_pool(fourth_stage_output).squeeze(dim=[2,3])
            pooled_mid_output = self.avg_pool(mid_output).squeeze(dim=[2,3])
            if self.do_classifier_free_guidance:
                pooled_mid_output_text, pooled_mid_output_ucond = pooled_mid_output.chunk(2, dim=0)
                pooled_mid_output = pooled_mid_output_ucond + self.cfg.guidance_scale * (pooled_mid_output_text - pooled_mid_output_ucond)

                if self.cfg.multi_scale_cfg:
                    pooled_first_stage_output_text, pooled_first_stage_output_ucond = pooled_first_stage_output.chunk(2, dim=0)
                    pooled_first_stage_output = pooled_first_stage_output_ucond + self.cfg.guidance_scale * (pooled_first_stage_output_text - pooled_first_stage_output_ucond)

                    pooled_second_stage_output_text, pooled_second_stage_output_ucond = pooled_second_stage_output.chunk(2, dim=0)
                    pooled_second_stage_output = pooled_second_stage_output_ucond + self.cfg.guidance_scale * (pooled_second_stage_output_text - pooled_second_stage_output_ucond)

                    pooled_third_stage_output_text, pooled_third_stage_output_ucond = pooled_third_stage_output.chunk(2, dim=0)
                    pooled_third_stage_output = pooled_third_stage_output_ucond + self.cfg.guidance_scale * (pooled_third_stage_output_text - pooled_third_stage_output_ucond)

                    pooled_fourth_stage_output_text, pooled_fourth_stage_output_ucond = pooled_fourth_stage_output.chunk(2, dim=0)
                    pooled_fourth_stage_output = pooled_fourth_stage_output_ucond + self.cfg.guidance_scale * (pooled_fourth_stage_output_text - pooled_fourth_stage_output_ucond)
                else:
                    pooled_first_stage_output_text, pooled_first_stage_output_ucond = pooled_first_stage_output.chunk(2, dim=0)
                    pooled_first_stage_output = pooled_first_stage_output_text

                    pooled_second_stage_output_text, pooled_second_stage_output_ucond = pooled_second_stage_output.chunk(2, dim=0)
                    pooled_second_stage_output = pooled_second_stage_output_text

                    pooled_third_stage_output_text, pooled_third_stage_output_ucond = pooled_third_stage_output.chunk(2, dim=0)
                    pooled_third_stage_output = pooled_third_stage_output_text

                    pooled_fourth_stage_output_text, pooled_fourth_stage_output_ucond = pooled_fourth_stage_output.chunk(2, dim=0)
                    pooled_fourth_stage_output = pooled_fourth_stage_output_text

            concat_pooled_output = torch.cat([pooled_first_stage_output, pooled_second_stage_output, pooled_third_stage_output, pooled_fourth_stage_output, pooled_mid_output], dim=-1)
            image_features = self.visual_projection(concat_pooled_output)

        else:
            pooled_mid_output = self.avg_pool(mid_output).squeeze(dim=[2,3])
            if self.do_classifier_free_guidance:
                pooled_mid_output_text, pooled_mid_output_ucond = pooled_mid_output.chunk(2, dim=0)
                pooled_mid_output = pooled_mid_output_ucond + self.cfg.guidance_scale * (pooled_mid_output_text - pooled_mid_output_ucond)
            image_features = self.visual_projection(pooled_mid_output)
        
        return image_features
    
    def forward(self, text_inputs, image_inputs, time_cond, generator=None):
        n_p = text_inputs.shape[0]
        n_i = image_inputs.shape[0]
        outputs = ()
        
        encoder_hidden_states, text_features = self.get_text_features(text_inputs)
        outputs += text_features,

        if n_i == 2 * n_p:
            if self.do_classifier_free_guidance:
                encoder_hidden_states_text, encoder_hidden_states_ucond = encoder_hidden_states.chunk(2, dim=0)
                encoder_hidden_states = torch.cat([encoder_hidden_states_text] * 2 + [encoder_hidden_states_ucond] * 2, dim=0)
            else:
                encoder_hidden_states = torch.cat([encoder_hidden_states, encoder_hidden_states], dim=0)
        image_features = self.get_image_features(encoder_hidden_states, image_inputs, time_cond, generator=generator)
        outputs += image_features,

        return outputs

    def save(self, path):
        self.unet.save_pretrained(os.path.join(path, "unet"), safe_serialization=True)
        if not self.cfg.freeze_text_encoder:
            self.text_encoder.save_pretrained(os.path.join(path, "text_encoder"), safe_serialization=True)
        
        # save others
        state_dict = {
            'visual_projection': self.visual_projection.state_dict(),
            'text_projection': self.text_projection.state_dict(),
            'logit_scale': self.logit_scale.data.item()
        }
        torch.save(state_dict, os.path.join(path, "state_dict.pt"))
        logger.info(f"Save model to path {path} successfully")
        
    def load(self, path):
        self.unet = self.unet.from_pretrained(os.path.join(path, "unet"))
        logger.info(f"Loading unet weights from {os.path.join(path, 'unet')}")
        if not self.cfg.freeze_text_encoder:
            self.text_encoder = self.text_encoder.from_pretrained(os.path.join(path, "text_encoder"))
            logger.info(f"Loading text_encoder weights from {os.path.join(path, 'text_encoder')}")
            
        # load others
        state_dict = torch.load(os.path.join(path, "state_dict.pt"))
        self.visual_projection.load_state_dict(state_dict['visual_projection'])
        self.text_projection.load_state_dict(state_dict['text_projection'])
        self.logit_scale.data = torch.tensor(state_dict['logit_scale'])       
        logger.info(f"Loading projection and logit_scale weights from {os.path.join(path, 'state_dict.pt')}")
        
    
    def encode_prompt(self, prompt):
        text_inputs = self.tokenizer(
            prompt,
            padding="max_length",
            max_length=self.tokenizer.model_max_length,
            truncation=True,
            return_tensors="pt",
        )
        return text_inputs.input_ids
    
    def preprocess_image(self, images):
        if not isinstance(images, list):
            images = [images]

        image_inputs = []
        for image in images:
            if isinstance(image, dict):
                image = image["bytes"]
            if isinstance(image, bytes):
                image = Image.open(BytesIO(image))
            elif isinstance(image, str):
                image = Image.open(image)
            image = image.convert("RGB")
            image = self.val_transform(image)     
            image_inputs.append(image)   
        image_inputs = torch.stack(image_inputs, dim=0)
        return image_inputs
    
    
    def get_preference_scores(self, prompt, images, timesteps, generator=None):
        image_inputs = self.preprocess_image(images).to(self.vae.device, dtype=self.vae.dtype)
        text_inputs = self.encode_prompt(prompt).to(self.text_encoder.device)
        timesteps = torch.tensor([timesteps] * image_inputs.shape[0], dtype=torch.long).to(self.vae.device)
        
        with torch.no_grad():
            text_embs, image_embs = self.forward(text_inputs, image_inputs, timesteps, generator=generator)
            
            image_embs = image_embs / torch.norm(image_embs, dim=-1, keepdim=True)
            text_embs = text_embs / torch.norm(text_embs, dim=-1, keepdim=True)

            scores = self.logit_scale.exp() * (text_embs @ image_embs.T)[0]
            
            probs = torch.softmax(scores, dim=-1)
        
        return scores.cpu().tolist(), probs.cpu().tolist()