LLaVA (Large Language and Vision Assistant) Implementation in PyTorch

This implementation builds LLaVA from scratch. It includes a ResNet-50 vision encoder with a projection MLP, a Transformer-based language model, a cross-attention layer for vision-language fusion, and an enhanced tokenizer. The full LLaVA module combines these components for image-grounded text generation.

import torch
import torch.nn as nn
import torchvision.models as models
from PIL import Image
import torchvision.transforms as transforms
import warnings
warnings.filterwarnings("ignore")

class LanguageModel(nn.Module):
    def __init__(self, vocab_size=1000, hidden_dim=768, num_heads=8, num_layers=2):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, hidden_dim)
        self.transformer = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(
                d_model=hidden_dim,
                nhead=num_heads,
                batch_first=True,
                activation='gelu'  
            ),
            num_layers=num_layers
        )
        self.output = nn.Linear(hidden_dim, vocab_size)
        self.layer_norm = nn.LayerNorm(hidden_dim) 

    def forward(self, input_ids=None, inputs_embeds=None, attention_mask=None):
        if inputs_embeds is None:
            inputs_embeds = self.embedding(input_ids)
        
        inputs_embeds = self.layer_norm(inputs_embeds)  
        
        if attention_mask is not None:
            hidden_states = self.transformer(inputs_embeds, src_key_padding_mask=~attention_mask.bool())
        else:
            hidden_states = self.transformer(inputs_embeds)
            
        logits = self.output(hidden_states)
        
        return type('ModelOutput', (), {
            'logits': logits,
            'hidden_states': hidden_states
        })()

class VisionEncoder(nn.Module):
    def __init__(self, output_dim=768, dropout_rate=0.1):
        super().__init__()
        resnet = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
        self.backbone = nn.Sequential(*list(resnet.children())[:-2]) 

        # pooling and projection layers
        self.avg_pool = nn.AdaptiveAvgPool2d((1, 1))
        self.dropout = nn.Dropout(dropout_rate)
        self.projection = nn.Sequential(
            nn.Linear(2048, output_dim * 2),
            nn.GELU(),
            nn.Linear(output_dim * 2, output_dim)
        )

    def forward(self, images):
        features = self.backbone(images)
        pooled = self.avg_pool(features).squeeze(-1).squeeze(-1)
        pooled = self.dropout(pooled)
        return self.projection(pooled)

class CrossAttention(nn.Module):
    def __init__(self, dim=768, num_heads=12, dropout=0.1):
        super().__init__()
        self.num_heads = num_heads
        self.scale = (dim // num_heads) ** -0.5
        
        self.q_proj = nn.Linear(dim, dim)
        self.k_proj = nn.Linear(dim, dim)
        self.v_proj = nn.Linear(dim, dim)
        self.out_proj = nn.Linear(dim, dim)
        
        self.dropout = nn.Dropout(dropout)
        self.layer_norm = nn.LayerNorm(dim)

    def forward(self, x, context):
        residual = x
        x = self.layer_norm(x)
        context = self.layer_norm(context)

        B, N, C = x.shape
        q = self.q_proj(x).reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
        k = self.k_proj(context).reshape(B, -1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
        v = self.v_proj(context).reshape(B, -1, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
        
        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = self.dropout(attn.softmax(dim=-1))

        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.out_proj(x)

        return x + residual 

class EnhancedTokenizer:
    def __init__(self, vocab_size=1000):
        self.vocab_size = vocab_size
        self.pad_token_id = vocab_size - 1
        self.eos_token_id = vocab_size - 2
        self.unk_token_id = vocab_size - 3

    def _tokenize_word(self, word):
        # subword tokenization
        if len(word) <= 3:
            return [ord(c) % (self.vocab_size-3) for c in word]
        return [ord(c) % (self.vocab_size-3) for c in word[:3]] + \
               [self.unk_token_id] + \
               [ord(c) % (self.vocab_size-3) for c in word[-3:]]

    def __call__(self, text, return_tensors=None, padding=None, max_length=None):
        if isinstance(text, str):
            text = [text]

        # Tokenize words
        token_ids = []
        for t in text:
            words = t.split()
            ids = []
            for word in words:
                ids.extend(self._tokenize_word(word))
            token_ids.append(ids)

        if max_length:
            token_ids = [ids[:max_length-1] + [self.eos_token_id] for ids in token_ids]
        
        if padding:
            max_len = max(len(ids) for ids in token_ids)
            attention_mask = [[1] * len(ids) + [0] * (max_len - len(ids)) for ids in token_ids]
            token_ids = [ids + [self.pad_token_id] * (max_len - len(ids)) for ids in token_ids]
        else:
            attention_mask = [[1] * len(ids) for ids in token_ids]

        if return_tensors == "pt":
            return type('TokenizerOutput', (), {
                'input_ids': torch.tensor(token_ids),
                'attention_mask': torch.tensor(attention_mask)
            })()
        
        return {'input_ids': token_ids, 'attention_mask': attention_mask}

class LLaVA(nn.Module):
    def __init__(self, vision_dim=768, vocab_size=1000):
        super().__init__()
        self.vision_encoder = VisionEncoder(output_dim=vision_dim)
        self.language_model = LanguageModel(vocab_size=vocab_size, hidden_dim=vision_dim)
        self.cross_attention = CrossAttention(dim=vision_dim)

        self.image_transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406],
                std=[0.229, 0.224, 0.225]
            )
        ])

        self.tokenizer = EnhancedTokenizer(vocab_size=vocab_size)

    def preprocess_image(self, image):
        if isinstance(image, Image.Image):
            try:
                return self.image_transform(image).unsqueeze(0)
            except Exception as e:
                raise ValueError(f"Error preprocessing image: {str(e)}")
        return image

    def forward(self, image, text):
        image = self.preprocess_image(image)
        vision_features = self.vision_encoder(image)

        # Process text
        text_inputs = self.tokenizer(text, return_tensors="pt", padding=True)
        lang_features = self.language_model.embedding(text_inputs.input_ids)
        
        # Fuse features
        fused_features = self.cross_attention(
            lang_features,
            vision_features.unsqueeze(1)
        )

        # Generate output
        outputs = self.language_model(
            inputs_embeds=fused_features,
            attention_mask=text_inputs.attention_mask
        )
        
        return outputs

def test_llava():
    print("Initializing LLaVA model...")
    model = LLaVA()
    model.eval() 

    print("\nRunning test cases...")

    # image-text processing
    test_image = Image.new('RGB', (224, 224))
    test_prompt = "What can you see in this image?"
    try:
        with torch.no_grad():
            outputs = model(test_image, test_prompt)
        print("✓ Basic processing passed")
        print(f"  Output logits shape: {outputs.logits.shape}")
    except Exception as e:
        print(f"✗ failed: {str(e)}")

    # Vision encoder output shape and values
    try:
        image_tensor = model.preprocess_image(test_image)
        vision_features = model.vision_encoder(image_tensor)
        assert vision_features.shape == (1, 768), f"Incorrect shape: {vision_features.shape}"
        assert not torch.isnan(vision_features).any(), "NaN values detected"
        print("✓ Vision encoder checks passed")
    except Exception as e:
        print(f"✗ failed: {str(e)}")

    # Tokenizer functionality
    try:
        test_texts = ["Hello world", "Testing multiple sentences"]
        tokens = model.tokenizer(test_texts, padding=True, return_tensors="pt")
        assert tokens.input_ids.shape[0] == 2, "Batch size mismatch"
        assert tokens.attention_mask.shape == tokens.input_ids.shape, "Attention mask shape mismatch"
        print("✓ Tokenizer checks passed")
    except Exception as e:
        print(f"✗ failed: {str(e)}")

    # Cross-attention mechanism
    try:
        batch_size, seq_len = 2, 10
        dummy_lang = torch.randn(batch_size, seq_len, 768)
        dummy_vision = torch.randn(batch_size, 768)
        fused = model.cross_attention(dummy_lang, dummy_vision.unsqueeze(1))
        assert fused.shape == (batch_size, seq_len, 768), f"Incorrect shape: {fused.shape}"
        assert not torch.isnan(fused).any(), "NaN values detected"
        print("✓ Cross-attention checks passed")
    except Exception as e:
        print(f"✗ failed: {str(e)}")

    # batched input
    try:
        batch_prompts = ["First prompt", "Second prompt"]
        batch_images = [Image.new('RGB', (224, 224)) for _ in range(2)]
        
        with torch.no_grad():
            for img, prompt in zip(batch_images, batch_prompts):
                outputs = model(img, prompt)
                assert outputs.logits.shape[1] == len(model.tokenizer(prompt)['input_ids'][0])
        print("✓ Batch processing passed")
    except Exception as e:
        print(f"✗ failed: {str(e)}")

if __name__ == "__main__":
    test_llava()