llava/model/frozen_llava_extractor.py

"""
Frozen LLaVA Feature Extractor

Loads a pretrained LLaVA model (with LoRA) and extracts hidden states before the language model head.
The entire model is frozen - no gradients computed.
Hidden states serve as a consistent "codebook" for guiding InstructPix2Pix.
"""

import torch
import torch.nn as nn
import os
from typing import Optional, Tuple

from llava.model.builder import load_pretrained_model
from llava.mm_utils import get_model_name_from_path, tokenizer_image_token
from llava.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN, DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN
from llava.conversation import conv_templates
from peft import PeftModel


class FrozenLLaVAExtractor(nn.Module):
    """
    Frozen LLaVA model for extracting hidden states as conditioning signals.
    
    Workflow:
    1. Load pretrained LLaVA + LoRA weights
    2. Freeze all parameters
    3. Extract hidden states before lm_head (shape: B, L, 4096)
    4. These act as consistent semantic "codebook" for diffusion guidance
    
    Args:
        llava_base_path: Path to base LLaVA model (e.g., "llava-1.5-7b-hf")
        llava_lora_path: Path to finetuned LoRA weights (e.g., "llava-miragehd-prior-bbox")
        device: Device to load model on
        load_8bit: Whether to load in 8-bit (for memory efficiency)
        load_4bit: Whether to load in 4-bit
    """
    def __init__(
        self,
        llava_base_path: str,
        llava_lora_path: Optional[str] = None,
        device: str = "cuda",
        load_8bit: bool = False,
        load_4bit: bool = False,
        merge_lora: bool = True,
    ):
        super().__init__()
        
        self.device = device
        self.llava_base_path = llava_base_path
        self.llava_lora_path = llava_lora_path
        
        print(f"Loading frozen LLaVA from: {llava_base_path}")
        if llava_lora_path:
            print(f"  + LoRA weights from: {llava_lora_path}")
        
        # Load base model
        # IMPORTANT: Load directly on the target CUDA device to avoid multi-GPU splits
        model_name = get_model_name_from_path(llava_base_path)
        target_device_map = {"": device} if isinstance(device, str) and device.startswith("cuda") else None
        self.tokenizer, self.model, self.image_processor, self.context_len = load_pretrained_model(
            llava_base_path,
            None,
            model_name,
            load_8bit=load_8bit,
            load_4bit=load_4bit,
            device_map=target_device_map if target_device_map is not None else device,
            device=device
        )
        
        # Load LoRA weights if provided
        if llava_lora_path:
            print("Loading LoRA adapter...")
            self.model = PeftModel.from_pretrained(self.model, llava_lora_path)
            
            # Load non-LoRA trainables (mm_projector weights)
            non_lora_path = os.path.join(llava_lora_path, 'non_lora_trainables.bin')
            if os.path.exists(non_lora_path):
                print("Loading mm_projector weights...")
                non_lora_weights = torch.load(non_lora_path, map_location='cpu')
                # Convert to fp16 if needed
                if not load_8bit and not load_4bit:
                    non_lora_weights = {k: v.to(torch.float16) for k, v in non_lora_weights.items()}
                self.model.load_state_dict(non_lora_weights, strict=False)
            else:
                print("Warning: non_lora_trainables.bin not found, mm_projector may not be updated")

            # Optionally merge LoRA to avoid runtime PEFT wrappers on vision/LLM
            if merge_lora:
                try:
                    print("Merging LoRA adapters into base model and unloading...")
                    self.model = self.model.merge_and_unload()
                except Exception as e:
                    print(f"Warning: merge_and_unload failed ({e}); continuing with PEFT wrappers.")
        
        # Freeze everything
        # Ensure model is on the target device (safety in multi-GPU runs)
        try:
            if isinstance(self.device, str) and self.device.startswith('cuda') and torch.cuda.is_available():
                # Set current device context to reduce cuBLAS init issues on some systems
                try:
                    local_idx = int(self.device.split(':')[1]) if ':' in self.device else 0
                    torch.cuda.set_device(local_idx)
                except Exception:
                    pass
            self.model.to(self.device)
        except Exception:
            # Some models already sharded by device_map; skip explicit .to()
            pass

        self.model.eval()
        for param in self.model.parameters():
            param.requires_grad = False
        
        print("✓ LLaVA model loaded and frozen")
        
    def prepare_inputs(
        self,
        images,  # Either list of PIL images or torch.Tensor
        prompts: list,
        conv_mode: str = "llava_v1"
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Prepare inputs for LLaVA forward pass.
        
        Args:
            images: Either list of PIL images or image tensors (B, 3, H, W)
            prompts: List of text prompts (B,)
            conv_mode: Conversation template mode
        
        Returns:
            input_ids: Tokenized inputs (B, L)
            images: Processed images (B, 3, 336, 336)
        """
        batch_size = len(prompts)
        
        # Process images: convert PIL to tensors if needed
        if isinstance(images, list):
            # List of PIL images - process each one
            from PIL import Image
            processed_images = []
            for img in images:
                if isinstance(img, Image.Image):
                    # Preprocess PIL image (LLaVA's CLIP processor)
                    img_tensor = self.image_processor.preprocess(img, return_tensors='pt')['pixel_values'][0]
                    processed_images.append(img_tensor)
                else:
                    # Already a tensor
                    processed_images.append(img)
            images = torch.stack(processed_images).to(self.device)
        else:
            # Already a batched tensor
            if images.device != self.device:
                images = images.to(self.device)
        
        # Tokenize prompts with image tokens
        input_ids_list = []
        for prompt in prompts:
            # Add image token
            if self.model.config.mm_use_im_start_end:
                prompt = DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + '\n' + prompt
            else:
                prompt = DEFAULT_IMAGE_TOKEN + '\n' + prompt
            
            # Tokenize
            input_ids = tokenizer_image_token(
                prompt, 
                self.tokenizer, 
                IMAGE_TOKEN_INDEX, 
                return_tensors='pt'
            )
            input_ids_list.append(input_ids)
        
        # Pad to same length
        max_len = max(ids.size(0) for ids in input_ids_list)
        input_ids_padded = []
        for ids in input_ids_list:
            padding = torch.full((max_len - ids.size(0),), self.tokenizer.pad_token_id, dtype=ids.dtype)
            input_ids_padded.append(torch.cat([ids, padding]))
        
        input_ids = torch.stack(input_ids_padded).to(self.device)
        
        return input_ids, images
    
    @torch.no_grad()
    def extract_hidden_states(
        self,
        images,  # Either list of PIL images or torch.Tensor
        prompts: list,
    ) -> torch.Tensor:
        """
        Extract hidden states from LLaVA model (before lm_head).
        
        Args:
            images: Either list of PIL images or image tensors (B, 3, H, W)
            prompts: Text prompts (B,)
        
        Returns:
            hidden_states: Hidden states before lm_head, shape (B, L, 4096)
        """
        # Prepare inputs (handles PIL->tensor conversion if needed)
        input_ids, image_tensors = self.prepare_inputs(images, prompts)
        
        # Ensure FP16 for efficiency
        if image_tensors.device.type == "cpu":
            # Avoid half on CPU (causes slow_conv2d_cpu not implemented for Half)
            pass
        else:
            if image_tensors.dtype != torch.float16:
                image_tensors = image_tensors.half()
        
        # Forward through LLaVA; use top-level model so it accepts `images`
        # hidden_states are from the base model (before lm_head)
        def _forward_with(images_tensor: torch.Tensor):
            return self.model(
                input_ids=input_ids,
                images=images_tensor,
                return_dict=True,
                output_hidden_states=True
            )

        try:
            images_fwd = image_tensors
            if images_fwd.device.type == 'cuda':
                images_fwd = images_fwd.contiguous()
            outputs = _forward_with(images_fwd)
        except RuntimeError as e:
            msg = str(e)
            # Fallback: retry in float32 on the same device for broader kernel support
            if 'GET was unable to find an engine' in msg or 'CUBLAS_STATUS_NOT_INITIALIZED' in msg:
                try:
                    vt = None
                    try:
                        vt = self.model.get_vision_tower()
                    except Exception:
                        vt = None
                    if vt is not None:
                        vt.to(device=self.device, dtype=torch.float32)
                    images_f32 = image_tensors.to(device=self.device, dtype=torch.float32).contiguous()
                    outputs = _forward_with(images_f32)
                except Exception:
                    raise
            else:
                raise
        
        # Get last hidden states (before lm_head)
        hidden_states = outputs.hidden_states[-1]  # Shape: (B, L, 4096)
        
        return hidden_states
    
    def forward(self, images, prompts: list) -> torch.Tensor:
        """
        Forward pass - extract hidden states.
        
        Args:
            images: Either list of PIL images or image tensors (B, 3, H, W)
            prompts: Text prompts (B,)
        
        Returns:
            hidden_states: Hidden states, shape (B, L, 4096)
        """
        return self.extract_hidden_states(images, prompts)


def test_extractor():
    """Test the frozen LLaVA extractor"""
    import torch
    from PIL import Image
    import requests
    from io import BytesIO
    
    print("Testing FrozenLLaVAExtractor...")
    
    # Paths (adjust as needed)
    llava_base = "checkpoints/llava-1.5-7b-hf"
    llava_lora = "checkpoints/llava-miragehd-prior-bbox"
    
    # Create extractor
    extractor = FrozenLLaVAExtractor(
        llava_base_path=llava_base,
        llava_lora_path=llava_lora,
        device="cuda"
    )
    
    # Test with dummy data
    batch_size = 2
    prompts = [
        "How would this RGB scene appear in long-wave thermal infrared spectrum",
        "Describe the thermal characteristics of this scene"
    ]
    
    # Create dummy images
    dummy_images = torch.randn(batch_size, 3, 336, 336).cuda()
    
    # Extract hidden states
    hidden_states = extractor(dummy_images, prompts)
    print(f"Hidden states shape: {hidden_states.shape}")
    print(f"Expected: (batch_size, variable_length, 4096)")
    print(f"✓ Frozen LLaVA extractor test passed!")
    
    # Verify frozen
    assert not any(p.requires_grad for p in extractor.parameters()), "Model should be frozen!"
    print("✓ All parameters frozen")


if __name__ == "__main__":
    test_extractor()