Reward_sd15_idealized/models/reward_model.py

"""
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)