""" LRM Reward Model Wrapper Loads LRM weights from HuggingFace and provides interface for computing preference scores on noisy latents. """ import torch from torch import nn from diffusers import AutoencoderKL, DDPMScheduler from transformers import CLIPTextModel, CLIPTokenizer from huggingface_hub import hf_hub_download import os from .unet_2d_condition_reward import UNet2DConditionModel def _offline_mode_enabled() -> bool: return os.getenv("HF_HUB_OFFLINE", "0").strip().lower() in {"1", "true", "yes", "on"} def _get_cache_dir() -> str | None: return os.getenv("HF_HUB_CACHE") or os.getenv("HUGGINGFACE_HUB_CACHE") def _hf_pretrained_kwargs() -> dict: kwargs = {"local_files_only": _offline_mode_enabled()} cache_dir = _get_cache_dir() if cache_dir: kwargs["cache_dir"] = cache_dir return kwargs class LRMRewardModel(nn.Module): """ Latent Reward Model (LRM) for SD1.5 This model computes preference scores directly on noisy latent images at any timestep. It uses features from the U-Net and text encoder to predict how well an image aligns with the prompt at different noise levels. Args: pretrained_model_name_or_path: Base SD model path (e.g., 'runwayml/stable-diffusion-v1-5') lrm_model_path: Path to LRM checkpoint from HuggingFace (e.g., 'casiatao/LRM') clip_model_path: Path to CLIP checkpoint for text projection initialization guidance_scale: Classifier-free guidance scale (default: 7.5) device: Device to load model on """ def __init__( self, pretrained_model_name_or_path='runwayml/stable-diffusion-v1-5', lrm_model_path=None, clip_model_path='openai/clip-vit-large-patch14', 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 base models from {pretrained_model_name_or_path}...") pretrained_kwargs = _hf_pretrained_kwargs() # Load tokenizer and text encoder self.tokenizer = CLIPTokenizer.from_pretrained( pretrained_model_name_or_path, subfolder="tokenizer", **pretrained_kwargs, ) self.text_encoder = CLIPTextModel.from_pretrained( pretrained_model_name_or_path, subfolder="text_encoder", **pretrained_kwargs, ).to(device) # Load VAE (frozen, only needed for preprocessing if using images) self.vae = AutoencoderKL.from_pretrained( pretrained_model_name_or_path, subfolder="vae", **pretrained_kwargs, ).to(device) self.vae.requires_grad_(False) # Load scheduler self.scheduler = DDPMScheduler.from_pretrained( pretrained_model_name_or_path, subfolder="scheduler", **pretrained_kwargs, ) # Load U-Net with custom reward architecture print("Loading custom U-Net for reward prediction...") self.unet = UNet2DConditionModel.from_pretrained( pretrained_model_name_or_path, subfolder="unet", **pretrained_kwargs, ).to(device) # Global pooling layer self.avg_pool = nn.AdaptiveAvgPool2d((1, 1)) # Projection layers # Multi-scale: concatenates features from 4 down blocks + mid block = 4800 dims vision_embed_dim = 4800 if self.multi_scale else 1280 text_embed_dim = 768 projection_dim = 768 self.visual_projection = nn.Linear(vision_embed_dim, projection_dim, bias=False).to(device) self.text_projection = nn.Linear(text_embed_dim, projection_dim, bias=False).to(device) # Initialize text projection from CLIP print(f"Loading CLIP text projection from {clip_model_path}...") try: # Try loading from local path first if os.path.exists(clip_model_path): clip_ckpt = torch.load(clip_model_path, map_location='cpu') else: # Download from HuggingFace clip_ckpt_path = hf_hub_download( repo_id=clip_model_path, filename="pytorch_model.bin", **_hf_pretrained_kwargs(), ) clip_ckpt = torch.load(clip_ckpt_path, map_location='cpu') self.text_projection.weight.data = clip_ckpt['text_projection.weight'].contiguous().to(device) print("✓ Loaded CLIP text projection weights") except Exception as e: print(f"Warning: Could not load CLIP weights: {e}") print("Initializing text projection randomly") nn.init.normal_(self.text_projection.weight, std=0.02) # Initialize visual projection nn.init.normal_(self.visual_projection.weight, std=0.02) # Logit scale (temperature parameter) self.logit_scale = nn.Parameter(torch.ones([]) * 2.6592).to(device) # Setup classifier-free guidance self.do_classifier_free_guidance = self.guidance_scale > 1.0 if self.do_classifier_free_guidance: self.neg_prompt_ids = self.tokenizer( [""], return_tensors="pt", padding="max_length", truncation=True, max_length=self.tokenizer.model_max_length, ).input_ids.to(device) # Load fine-tuned LRM weights if provided if lrm_model_path: self.load_lrm_weights(lrm_model_path) print("✓ LRM Reward Model initialized successfully!") def load_lrm_weights(self, model_path): """ Load fine-tuned LRM weights from HuggingFace or local path Expected structure: - unet/ (directory with U-Net weights) - text_encoder/ (optional, directory with text encoder weights) - state_dict.pt (visual_projection, text_projection, logit_scale) """ print(f"\nLoading LRM weights from {model_path}...") try: # Check if it's a HuggingFace model or local path if not os.path.exists(model_path): # Try to download from HuggingFace print(f"Downloading from HuggingFace: {model_path}") # For HF models, we need to download the entire repo from huggingface_hub import snapshot_download model_path = snapshot_download(repo_id=model_path, **_hf_pretrained_kwargs()) # Load U-Net weights unet_path = os.path.join(model_path, "lrm_sd15", "unet") if os.path.exists(unet_path): self.unet = UNet2DConditionModel.from_pretrained(unet_path, **_hf_pretrained_kwargs()).to(self.device) print(f"✓ Loaded U-Net weights from {unet_path}") else: print(f"Warning: U-Net path not found: {unet_path}") # Load text encoder weights (optional) text_encoder_path = os.path.join(model_path, "lrm_sd15", "text_encoder") if os.path.exists(text_encoder_path): self.text_encoder = CLIPTextModel.from_pretrained(text_encoder_path, **_hf_pretrained_kwargs()).to(self.device) print(f"✓ Loaded text encoder weights from {text_encoder_path}") # Load projection layers and logit scale state_dict_path = os.path.join(model_path, "lrm_sd15", "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.text_projection.load_state_dict(state_dict['text_projection']) # Move projection layers to device self.visual_projection = self.visual_projection.to(self.device) self.text_projection = self.text_projection.to(self.device) logit_scale_val = state_dict['logit_scale'] if isinstance(logit_scale_val, torch.Tensor): self.logit_scale.data = logit_scale_val.to(self.device) else: self.logit_scale.data = torch.tensor(logit_scale_val).to(self.device) print(f"✓ Loaded projection layers and logit_scale from {state_dict_path}") else: print(f"Warning: state_dict.pt not found: {state_dict_path}") print("✓ Successfully loaded all LRM weights!") except Exception as e: print(f"Error loading LRM weights: {e}") print("Continuing with base model weights...") def encode_prompt(self, prompt): """Tokenize text prompt""" if isinstance(prompt, str): prompt = [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.to(self.device) def get_text_features(self, text_input_ids): """ Extract text features from prompt Returns: (encoder_hidden_states, text_features) """ if self.do_classifier_free_guidance: # Concatenate conditional and unconditional prompts text_input_ids = torch.cat([ text_input_ids, self.neg_prompt_ids.repeat(text_input_ids.shape[0], 1) ], dim=0) outputs = self.text_encoder(text_input_ids, return_dict=False) encoder_hidden_states = outputs[0] # Sequence of hidden states pooled_output = outputs[1] # Pooled output (last token) 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, noisy_latents, timesteps): """ Extract visual features from noisy latents using U-Net Args: encoder_hidden_states: Text conditioning from CLIP noisy_latents: Noisy latent images [B, C, H, W] timesteps: Denoising timesteps [B] Returns: image_features: Visual embeddings [B, projection_dim] """ if self.do_classifier_free_guidance: noisy_latents = torch.cat([noisy_latents] * 2, dim=0) timesteps = torch.cat([timesteps] * 2, dim=0) # Forward through U-Net (only down blocks + mid block, no up blocks) mid_output, down_block_res_samples = self.unet( noisy_latents, timesteps, encoder_hidden_states=encoder_hidden_states, return_dict=False, use_up_blocks=False ) if self.multi_scale: # Extract multi-scale features from down blocks # Indices correspond to: [320, 64, 64], [640, 32, 32], [1280, 16, 16], [1280, 8, 8] first_stage_output = down_block_res_samples[2] # 320 channels second_stage_output = down_block_res_samples[5] # 640 channels third_stage_output = down_block_res_samples[8] # 1280 channels fourth_stage_output = down_block_res_samples[11] # 1280 channels # Global average pooling pooled_first = self.avg_pool(first_stage_output).squeeze(dim=[2, 3]) pooled_second = self.avg_pool(second_stage_output).squeeze(dim=[2, 3]) pooled_third = self.avg_pool(third_stage_output).squeeze(dim=[2, 3]) pooled_fourth = self.avg_pool(fourth_stage_output).squeeze(dim=[2, 3]) pooled_mid = self.avg_pool(mid_output).squeeze(dim=[2, 3]) # Apply VFE (Visual Feature Enhancement) on mid block if self.do_classifier_free_guidance: pooled_mid_text, pooled_mid_ucond = pooled_mid.chunk(2, dim=0) pooled_mid = pooled_mid_ucond + self.guidance_scale * (pooled_mid_text - pooled_mid_ucond) # For other blocks, optionally apply CFG or just use conditional branch if self.multi_scale_cfg: pooled_first_text, pooled_first_ucond = pooled_first.chunk(2, dim=0) pooled_first = pooled_first_ucond + self.guidance_scale * (pooled_first_text - pooled_first_ucond) pooled_second_text, pooled_second_ucond = pooled_second.chunk(2, dim=0) pooled_second = pooled_second_ucond + self.guidance_scale * (pooled_second_text - pooled_second_ucond) pooled_third_text, pooled_third_ucond = pooled_third.chunk(2, dim=0) pooled_third = pooled_third_ucond + self.guidance_scale * (pooled_third_text - pooled_third_ucond) pooled_fourth_text, pooled_fourth_ucond = pooled_fourth.chunk(2, dim=0) pooled_fourth = pooled_fourth_ucond + self.guidance_scale * (pooled_fourth_text - pooled_fourth_ucond) else: # Use only conditional (text-conditioned) branch pooled_first, _ = pooled_first.chunk(2, dim=0) pooled_second, _ = pooled_second.chunk(2, dim=0) pooled_third, _ = pooled_third.chunk(2, dim=0) pooled_fourth, _ = pooled_fourth.chunk(2, dim=0) # Concatenate all scales: 320 + 640 + 1280 + 1280 + 1280 = 4800 concat_pooled = torch.cat([ pooled_first, pooled_second, pooled_third, pooled_fourth, pooled_mid ], dim=-1) image_features = self.visual_projection(concat_pooled) else: # Single scale (mid block only) pooled_mid = self.avg_pool(mid_output).squeeze(dim=[2, 3]) if self.do_classifier_free_guidance: pooled_mid_text, pooled_mid_ucond = pooled_mid.chunk(2, dim=0) pooled_mid = pooled_mid_ucond + self.guidance_scale * (pooled_mid_text - pooled_mid_ucond) image_features = self.visual_projection(pooled_mid) return image_features def get_reward_score(self, noisy_latents, prompt, timesteps, enable_grad=False): """ Compute preference score for noisy latents at given timesteps Args: noisy_latents: Noisy latent images [B, C, H, W] prompt: Text prompt(s) (string or list of strings) timesteps: Denoising timesteps [B] or scalar enable_grad: If True, allows gradient computation (for gradient ascent) Returns: scores: Preference scores [B] """ def _compute(): # Ensure inputs are on correct device latents = noisy_latents.to(self.device, dtype=self.unet.dtype) # Handle timesteps if isinstance(timesteps, int): ts = torch.tensor([timesteps] * latents.shape[0]) else: ts = timesteps ts = ts.to(self.device) # Encode prompt text_input_ids = self.encode_prompt(prompt) # Get text and image features encoder_hidden_states, text_features = self.get_text_features(text_input_ids) image_features = self.get_image_features(encoder_hidden_states, latents, ts) # Normalize features image_features = image_features / torch.norm(image_features, dim=-1, keepdim=True) text_features = text_features / torch.norm(text_features, dim=-1, keepdim=True) # Compute similarity scores scores = self.logit_scale.exp() * (text_features @ image_features.T)[0] scores = torch.sigmoid(scores) # Scale to [0, 1] return scores # return scores # If enable_grad is True, compute with gradients; otherwise use no_grad if enable_grad: return _compute() else: with torch.no_grad(): return _compute() def forward(self, noisy_latents, prompt, timesteps): """Alias for get_reward_score for nn.Module compatibility""" return self.get_reward_score(noisy_latents, prompt, timesteps)