""" 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) # Load scheduler 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) #print(torch.sigmoid(scores)) return torch.sigmoid(scores) #return scores