| """ |
| LRM Reward Model Wrapper for SDXL |
| """ |
|
|
| import torch |
| from torch import nn |
| from diffusers import AutoencoderKL, DDPMScheduler |
| from transformers import CLIPTextModel, CLIPTextModelWithProjection, CLIPTokenizer |
| import os |
|
|
| from .unet_2d_condition_reward import UNet2DConditionModel |
|
|
| class LRMRewardModelXL(nn.Module): |
| def __init__( |
| self, |
| pretrained_model_name_or_path='stabilityai/stable-diffusion-xl-base-1.0', |
| lrm_model_path=None, |
| guidance_scale=7.5, |
| device='cuda' |
| ): |
| super().__init__() |
| self.device = device |
| self.guidance_scale = guidance_scale |
| self.multi_scale = True |
| self.multi_scale_cfg = False |
| |
| print(f"Loading SDXL base models from {pretrained_model_name_or_path}...") |
| |
| self.tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder="tokenizer") |
| self.tokenizer_2 = CLIPTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder="tokenizer_2") |
| |
| self.text_encoder = CLIPTextModel.from_pretrained(pretrained_model_name_or_path, subfolder="text_encoder").to(device) |
| self.text_encoder_2 = CLIPTextModelWithProjection.from_pretrained(pretrained_model_name_or_path, subfolder="text_encoder_2").to(device) |
| self.unet = UNet2DConditionModel.from_pretrained(pretrained_model_name_or_path, subfolder="unet").to(device) |
|
|
| |
| self.scheduler = DDPMScheduler.from_pretrained( |
| pretrained_model_name_or_path, |
| subfolder="scheduler" |
| ) |
| |
| self.avg_pool = nn.AdaptiveAvgPool2d((1, 1)) |
| |
| vision_embed_dim = 3520 if self.multi_scale else 1280 |
| projection_dim = 1280 |
| |
| self.visual_projection = nn.Linear(vision_embed_dim, projection_dim, bias=False).to(device) |
| nn.init.normal_(self.visual_projection.weight, std=0.02) |
| |
| self.logit_scale = nn.Parameter(torch.ones([]) * 2.6592).to(device) |
| self.do_classifier_free_guidance = self.guidance_scale > 1.0 |
| |
| if lrm_model_path: |
| self.load_lrm_weights(lrm_model_path) |
| |
| def load_lrm_weights(self, model_path): |
| """Load fine-tuned SDXL LRM weights.""" |
| print(f"\nLoading LRM XL weights from {model_path}...") |
| try: |
| if not os.path.exists(model_path): |
| from huggingface_hub import snapshot_download |
| model_path = snapshot_download(repo_id=model_path) |
| |
| unet_path = os.path.join(model_path, "lrm_sdxl", "unet") |
| if os.path.exists(unet_path): |
| self.unet = UNet2DConditionModel.from_pretrained(unet_path).to(self.device) |
| print(f"? Loaded U-Net weights from {unet_path}") |
| |
| text_encoder_path = os.path.join(model_path, "lrm_sdxl", "text_encoder") |
| if os.path.exists(text_encoder_path): |
| self.text_encoder = CLIPTextModel.from_pretrained(text_encoder_path).to(self.device) |
| print(f"? Loaded Text Encoder 1 weights") |
| |
| text_encoder_2_path = os.path.join(model_path, "lrm_sdxl", "text_encoder_2") |
| if os.path.exists(text_encoder_2_path): |
| self.text_encoder_2 = CLIPTextModelWithProjection.from_pretrained(text_encoder_2_path).to(self.device) |
| print(f"? Loaded Text Encoder 2 weights") |
| |
| state_dict_path = os.path.join(model_path, "lrm_sdxl", "state_dict.pt") |
| if os.path.exists(state_dict_path): |
| state_dict = torch.load(state_dict_path, map_location='cpu') |
| self.visual_projection.load_state_dict(state_dict['visual_projection']) |
| self.logit_scale.data = torch.tensor(state_dict['logit_scale']).to(self.device) |
| print(f"? Loaded projection layers") |
| |
| except Exception as e: |
| print(f"Error loading LRM weights: {e}") |
|
|
| def get_text_features(self, prompt, batch_size, latent_height=128, latent_width=128): |
| text_inputs_1 = self.tokenizer(prompt, padding="max_length", max_length=self.tokenizer.model_max_length, truncation=True, return_tensors="pt") |
| prompt_embeds_1 = self.text_encoder(text_inputs_1.input_ids.to(self.device), output_hidden_states=True) |
| |
| text_inputs_2 = self.tokenizer_2(prompt, padding="max_length", max_length=self.tokenizer_2.model_max_length, truncation=True, return_tensors="pt") |
| prompt_embeds_2 = self.text_encoder_2(text_inputs_2.input_ids.to(self.device), output_hidden_states=True) |
| |
| prompt_embeds = torch.cat([prompt_embeds_1.hidden_states[-2], prompt_embeds_2.hidden_states[-2]], dim=-1) |
| pooled_prompt_embeds = prompt_embeds_2[0] |
|
|
| if self.do_classifier_free_guidance: |
| uncond_inputs_1 = self.tokenizer([""] * batch_size, padding="max_length", max_length=self.tokenizer.model_max_length, return_tensors="pt") |
| uncond_embeds_1 = self.text_encoder(uncond_inputs_1.input_ids.to(self.device), output_hidden_states=True).hidden_states[-2] |
| |
| uncond_inputs_2 = self.tokenizer_2([""] * batch_size, padding="max_length", max_length=self.tokenizer_2.model_max_length, return_tensors="pt") |
| uncond_embeds_2 = self.text_encoder_2(uncond_inputs_2.input_ids.to(self.device), output_hidden_states=True) |
| |
| uncond_prompt_embeds = torch.cat([uncond_embeds_1, uncond_embeds_2.hidden_states[-2]], dim=-1) |
| prompt_embeds = torch.cat([uncond_prompt_embeds, prompt_embeds], dim=0) |
| pooled_prompt_embeds = torch.cat([uncond_embeds_2[0], pooled_prompt_embeds], dim=0) |
|
|
| h, w = latent_height * 8, latent_width * 8 |
| add_time_ids = torch.tensor([[h, w, 0, 0, h, w]], dtype=prompt_embeds.dtype).to(self.device) |
| add_time_ids = add_time_ids.repeat(batch_size, 1) |
| if self.do_classifier_free_guidance: |
| add_time_ids = torch.cat([add_time_ids, add_time_ids], dim=0) |
|
|
| text_reward_features = pooled_prompt_embeds.chunk(2, dim=0)[1] if self.do_classifier_free_guidance else pooled_prompt_embeds |
| |
| return prompt_embeds, pooled_prompt_embeds, add_time_ids, text_reward_features |
|
|
| def get_image_features(self, prompt_embeds, pooled_prompt_embeds, add_time_ids, noisy_latents, timesteps): |
| if self.do_classifier_free_guidance: |
| noisy_latents = torch.cat([noisy_latents] * 2, dim=0) |
| timesteps = torch.cat([timesteps] * 2, dim=0) |
|
|
| added_cond_kwargs = { |
| "text_embeds": pooled_prompt_embeds.to(dtype=self.unet.dtype), |
| "time_ids": add_time_ids.to(dtype=self.unet.dtype) |
| } |
| |
| mid_output, down_block_res_samples = self.unet( |
| noisy_latents.to(dtype=self.unet.dtype), |
| timesteps, |
| encoder_hidden_states=prompt_embeds.to(dtype=self.unet.dtype), |
| added_cond_kwargs=added_cond_kwargs, |
| return_dict=False, |
| use_up_blocks=False |
| ) |
|
|
| if self.multi_scale: |
| first_stage, second_stage, third_stage = down_block_res_samples[2], down_block_res_samples[5], down_block_res_samples[8] |
| pooled_1 = self.avg_pool(first_stage).squeeze(dim=[2, 3]) |
| pooled_2 = self.avg_pool(second_stage).squeeze(dim=[2, 3]) |
| pooled_3 = self.avg_pool(third_stage).squeeze(dim=[2, 3]) |
| pooled_mid = self.avg_pool(mid_output).squeeze(dim=[2, 3]) |
|
|
| if self.do_classifier_free_guidance: |
| p1_u, p1_t = pooled_1.chunk(2, dim=0) |
| p2_u, p2_t = pooled_2.chunk(2, dim=0) |
| p3_u, p3_t = pooled_3.chunk(2, dim=0) |
| pm_u, pm_t = pooled_mid.chunk(2, dim=0) |
| |
| pooled_1 = p1_u + self.guidance_scale * (p1_t - p1_u) if self.multi_scale_cfg else p1_t |
| pooled_2 = p2_u + self.guidance_scale * (p2_t - p2_u) if self.multi_scale_cfg else p2_t |
| pooled_3 = p3_u + self.guidance_scale * (p3_t - p3_u) if self.multi_scale_cfg else p3_t |
| pooled_mid = pm_u + self.guidance_scale * (pm_t - pm_u) |
| |
| concat_pooled = torch.cat([pooled_1, pooled_2, pooled_3, pooled_mid], dim=-1) |
| image_features = self.visual_projection(concat_pooled.to(self.visual_projection.weight.dtype)) |
| else: |
| pooled_mid = self.avg_pool(mid_output).squeeze(dim=[2, 3]) |
| if self.do_classifier_free_guidance: |
| pm_u, pm_t = pooled_mid.chunk(2, dim=0) |
| pooled_mid = pm_u + self.guidance_scale * (pm_t - pm_u) |
| image_features = self.visual_projection(pooled_mid.to(self.visual_projection.weight.dtype)) |
|
|
| return image_features |
|
|
| def get_reward_score(self, noisy_latents, prompt, timesteps, enable_grad=False): |
| with torch.set_grad_enabled(enable_grad): |
| latents = noisy_latents.to(self.device, dtype=self.unet.dtype) |
| batch_size = latents.shape[0] |
| ts = torch.tensor([timesteps] * batch_size).to(self.device) if isinstance(timesteps, int) else timesteps.to(self.device) |
| |
| prompt_list = [prompt] if isinstance(prompt, str) else prompt |
| prompt_embeds, pooled_prompt_embeds, add_time_ids, text_features = self.get_text_features( |
| prompt_list, batch_size, latent_height=latents.shape[2], latent_width=latents.shape[3] |
| ) |
| image_features = self.get_image_features(prompt_embeds, pooled_prompt_embeds, add_time_ids, latents, ts) |
| |
| image_features = image_features.to(torch.float32) |
| text_features = text_features.to(torch.float32) |
| |
| image_features = image_features / torch.norm(image_features, dim=-1, keepdim=True) |
| text_features = text_features / torch.norm(text_features, dim=-1, keepdim=True) |
| |
| scores = self.logit_scale.float().exp() * (text_features * image_features).sum(dim=-1) |
| |
| return torch.sigmoid(scores) |
| |
