init

README.md

@@ -0,0 +1,123 @@
+---
+datasets:
+  - HuggingFaceM4/WebSight
+base_model:
+  - Salesforce/codet5-base
+  - google/siglip2-base-patch16-512
+---
+
+# LiteVit5 - Image-to-HTML Model
+
+A lightweight transformer model combining SigLIP vision encoder with T5 seq2seq decoder for image-to-text generation tasks.
+
+## Model Architecture
+
+- **Vision Encoder**: SigLIP2 (frozen)
+- **Vision Processing**: Multi-view fusion
+- **Seq2Seq Decoder**: CodeT5-based decoder with language modeling head
+- **Input**: Images (5 views per sample - 4 quarter views + 1 full view)
+- **Output**: Generated HTML
+
+## Installation
+
+```bash
+uv add transformers torch accelerate
+```
+
+## Usage
+
+### Loading the Model
+
+```python
+from transformers import AutoModel, AutoTokenizer
+from transformers import SiglipProcessor
+
+# Load the model
+model = AutoModel.from_pretrained("LiteVit5/model", trust_remote_code=True)
+
+# Load tokenizer and processor
+tokenizer = AutoTokenizer.from_pretrained("Salesforce/codet5-base")
+processor = SiglipProcessor.from_pretrained("google/siglip2-base-patch16-512")
+```
+
+### Inference Example
+
+```python
+from PIL import Image
+import torch
+
+from transformers import AutoModel, AutoTokenizer
+from transformers import SiglipProcessor
+
+# Load the model
+model = AutoModel.from_pretrained("LiteVit5/model", trust_remote_code=True, device_map="auto")
+
+# Load tokenizer and processor
+tokenizer = AutoTokenizer.from_pretrained("Salesforce/codet5-base")
+processor = SiglipProcessor.from_pretrained("google/siglip2-base-patch16-512")
+
+# Preprocess image (split into 4 parts + full image = 5 views)
+def prepare_image(image_path: str, processor):
+    """
+    Prepare image with 5 views (4 quarters + full).
+
+    Args:
+        image_path: Path to the image file
+        processor: SigLIP processor
+
+    Returns:
+        Tensor of shape [5, 3, 512, 512]
+    """
+    image = Image.open(image_path).convert("RGB")
+
+    # Split into 4 quarters
+    width, height = image.size
+    quarters = [
+        image.crop((0, 0, width//2, height//2)),           # top-left
+        image.crop((width//2, 0, width, height//2)),       # top-right
+        image.crop((0, height//2, width//2, height)),      # bottom-left
+        image.crop((width//2, height//2, width, height)),  # bottom-right
+    ]
+
+    # Process all views
+    processed = [
+        processor(images=q, return_tensors="pt")["pixel_values"]
+        for q in quarters
+    ]
+    # Add full image
+    processed.append(
+        processor(images=image, return_tensors="pt")["pixel_values"]
+    )
+
+    pixel_values = torch.cat(processed, dim=0)
+    return pixel_values
+
+def generate_text(model, pixel_values, tokenizer, max_length=512):
+    """
+    Generate text from image.
+
+    Args:
+        model: LiteVit5 model
+        pixel_values: Preprocessed image tensor
+        tokenizer: Tokenizer for decoding
+        max_length: Maximum generation length
+
+    Returns:
+        Generated text string
+    """
+    with torch.no_grad():
+        output_ids = model.generate(pixel_values, max_length=max_length)
+
+    text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    return text
+
+device = next(model.parameters()).device
+
+# Process images
+pixel_values = prepare_image("./image_13.png", processor)
+pixel_values = pixel_values.to(device)
+print("\nGenerating HTML from image_13.png...")
+output_ids = model.generate(pixel_values, max_length=2024)
+text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+print(f"Generated: {text}")
+```

config.json

@@ -0,0 +1,14 @@
+{
+  "architectures": ["LiteVit5ForConditionalGeneration"],
+  "model_type": "litevit5",
+  "pad_token_id": 0,
+  "eos_token_id": 2,
+  "decoder_start_token_id": 0,
+  "torch_dtype": "float16",
+  "transformers_version": "4.57.3",
+  "auto_map": {
+    "AutoConfig": "config.LiteVit5Config",
+    "AutoModel": "modeling_litevit5.LiteVit5ForConditionalGeneration",
+    "AutoModelForSeq2SeqLM": "modeling_litevit5.LiteVit5ForConditionalGeneration"
+  }
+}

config.py

@@ -0,0 +1,23 @@
+from transformers import PretrainedConfig
+
+
+class LiteVit5Config(PretrainedConfig):
+    """
+    Configuration class for LiteVit5ForConditionalGeneration.
+    """
+    
+    model_type = "litevit5"
+    
+    def __init__(
+        self,
+        pad_token_id: int = 0,
+        eos_token_id: int = 1,
+        decoder_start_token_id: int = 0,
+        **kwargs
+    ):
+        super().__init__(
+            pad_token_id=pad_token_id,
+            eos_token_id=eos_token_id,
+            decoder_start_token_id=decoder_start_token_id,
+            **kwargs
+        )

modeling_litevit5.py

@@ -0,0 +1,285 @@
+import math
+from typing import Optional, Tuple
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel, AutoModelForSeq2SeqLM, SiglipVisionModel
+from transformers.modeling_outputs import Seq2SeqLMOutput
+from .config import LiteVit5Config
+
+
+class LiteVit5ForConditionalGeneration(PreTrainedModel):
+    """
+    LiteVit5 model for vision-to-text generation tasks.
+    Combines SigLIP vision encoder with T5 seq2seq decoder for image-to-text tasks.
+    """
+    
+    config_class = LiteVit5Config
+    base_model_prefix = "litevit5"
+    
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        
+        # Vision model (frozen)
+        self.vision_model = SiglipVisionModel.from_pretrained(
+            "google/siglip2-base-patch16-512",
+            dtype=torch.float16
+        )
+        self.vision_model.eval()
+        for param in self.vision_model.parameters():
+            param.requires_grad = False
+        
+        # Load seq2seq decoder and lm_head from CodeT5
+        seq2seq_model = AutoModelForSeq2SeqLM.from_pretrained(
+            "Salesforce/codet5-base",
+            dtype=torch.float16
+        )
+        self.seq2seq_decoder = seq2seq_model.decoder
+        self.seq2seq_lm_head = seq2seq_model.lm_head
+        self._shift_right = seq2seq_model._shift_right
+        
+        # Vision processing layers
+        self.downsampler = nn.Conv2d(768, 768, kernel_size=2, stride=2, bias=False, dtype=torch.float16)
+        self.fuse = nn.Linear(768 * 2, 768).half()
+        self.pos_embedding = nn.Parameter(torch.zeros(1, 1024, 768, dtype=torch.float16), requires_grad=True)
+        self.linear_projection = nn.Linear(768, 768).half()
+        
+        self.post_init()
+    
+    def get_encoder(self):
+        """Return the vision encoder for the model."""
+        return self.vision_model
+    
+    def get_decoder(self):
+        """Return the seq2seq decoder."""
+        return self.seq2seq_decoder
+    
+    def _encode_vision(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        """
+        Encode image inputs into vision features.
+        
+        Args:
+            pixel_values: Input images of shape [B*5, 3, 512, 512] (5 views per sample)
+        
+        Returns:
+            Encoded vision features of shape [B, 1024, 768]
+        """
+        # Ensure pixel_values are float16
+        pixel_values = pixel_values.half()
+        
+        batch_size = pixel_values.size(0) // 5
+        scale = 5  # Number of views (4 quarter views + 1 full view)
+        num_patches = 32
+        
+        # Get vision embeddings
+        with torch.no_grad():
+            vision_model_outputs = self.vision_model(pixel_values=pixel_values)
+            vision_hidden_states = vision_model_outputs.last_hidden_state  # [B*5, 1024, 768]
+        
+        # Reshape to separate views
+        vision_hidden_states = vision_hidden_states.view(batch_size, scale, *vision_hidden_states.shape[1:])  # [B, 5, 1024, 768]
+        
+        # Process quarter views
+        quarters = vision_hidden_states[:, :4]  # [B, 4, 1024, 768]
+        quarters = quarters.view(batch_size, 4, num_patches, num_patches, -1)  # [B, 4, 32, 32, 768]
+        
+        # Combine quarter views into full image
+        upper = torch.cat([quarters[:, 0], quarters[:, 1]], dim=2)  # [B, 32, 64, 768]
+        lower = torch.cat([quarters[:, 2], quarters[:, 3]], dim=2)  # [B, 32, 64, 768]
+        pooled_image = torch.cat([upper, lower], dim=1)  # [B, 64, 64, 768]
+        pooled_image = pooled_image.permute(0, 3, 1, 2)  # [B, 768, 64, 64]
+        
+        # Downsample
+        pooled32 = self.downsampler(pooled_image)  # [B, 768, 32, 32]
+        pooled_tok = pooled32.flatten(2).transpose(1, 2)  # [B, 1024, 768]
+        
+        # Full image features
+        full_image = vision_hidden_states[:, 4]  # [B, 1024, 768]
+        
+        # Fuse quarter and full views
+        concat = torch.cat([pooled_tok, full_image], dim=-1)  # [B, 1024, 1536]
+        fused = self.fuse(concat)  # [B, 1024, 768]
+        
+        # Add positional encoding and project
+        fused = fused + self.pos_embedding
+        vision_hidden_states = self.linear_projection(fused)  # [B, 1024, 768]
+        
+        return vision_hidden_states
+    
+    def forward(
+        self,
+        pixel_values: torch.Tensor,
+        input_ids: Optional[torch.LongTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        decoder_input_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Tuple] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        **kwargs
+    ) -> Seq2SeqLMOutput:
+        """
+        Forward pass for the model.
+        
+        Args:
+            pixel_values: Vision input images
+            input_ids: Decoder input token IDs
+            labels: Target token IDs for training
+            decoder_input_ids: Decoder input IDs (used during generation)
+            past_key_values: Cached key values for efficient generation
+            attention_mask: Attention mask for decoder inputs
+        
+        Returns:
+            Seq2SeqLMOutput with loss, logits, and generation-related outputs
+        """
+        # Encode images
+        encoder_hidden_states = self._encode_vision(pixel_values)
+        
+        # Prepare decoder input IDs
+        if decoder_input_ids is None and input_ids is None:
+            decoder_input_ids = self._get_decoder_start_token_id()
+            decoder_input_ids = torch.full(
+                (pixel_values.shape[0] // 5, 1),
+                decoder_input_ids,
+                dtype=torch.long,
+                device=pixel_values.device
+            )
+        
+        if decoder_input_ids is None and input_ids is not None:
+            decoder_input_ids = self._shift_right(input_ids)
+        
+        # Pass through decoder
+        decoder_outputs = self.seq2seq_decoder(
+            input_ids=decoder_input_ids,
+            encoder_hidden_states=encoder_hidden_states,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+        )
+        
+        sequence_output = decoder_outputs[0]
+        lm_logits = self.seq2seq_lm_head(sequence_output)
+        
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
+            labels = labels.to(lm_logits.device)
+            loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), labels.view(-1))
+        
+        return Seq2SeqLMOutput(
+            loss=loss,
+            logits=lm_logits,
+            past_key_values=decoder_outputs.past_key_values,
+            decoder_hidden_states=decoder_outputs.hidden_states,
+            decoder_attentions=decoder_outputs.attentions,
+            cross_attentions=decoder_outputs.cross_attentions,
+        )
+    
+    def prepare_inputs_for_generation(
+        self,
+        decoder_input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        use_cache=None,
+        encoder_outputs=None,
+        **kwargs
+    ):
+        """Prepare inputs for generation."""
+        # Cut decoder_input_ids if past is used
+        if past_key_values is not None:
+            decoder_input_ids = decoder_input_ids[:, -1:]
+
+        return {
+            "input_ids": None,  # encoder_outputs is already defined
+            "encoder_outputs": encoder_outputs,
+            "past_key_values": past_key_values,
+            "decoder_input_ids": decoder_input_ids,
+            "attention_mask": attention_mask,
+            "use_cache": use_cache,
+        }
+
+    def _prepare_encoder_decoder_kwargs_for_generation(
+        self, inputs_tensor: torch.Tensor, model_kwargs, model_input_name: Optional[str] = None
+    ):
+        """Encode pixel values to get encoder outputs."""
+        # Encode images if not already done
+        if "encoder_outputs" not in model_kwargs:
+            encoder_outputs = self._encode_vision(inputs_tensor)
+            model_kwargs["encoder_outputs"] = (encoder_outputs,)
+        
+        return model_kwargs
+
+    def generate(
+        self,
+        pixel_values: torch.Tensor,
+        max_length: int = 1024,
+        num_beams: int = 1,
+        temperature: float = 1.0,
+        do_sample: bool = False,
+        **kwargs
+    ) -> torch.LongTensor:
+        """
+        Generate text from image inputs.
+        
+        Args:
+            pixel_values: Input images [B*5, 3, 512, 512]
+            max_length: Maximum generation length
+            num_beams: Number of beams for beam search (1 = greedy) TODO: Not implemented
+            temperature: Sampling temperature
+            do_sample: Whether to use sampling
+        
+        Returns:
+            Generated token sequences
+        """
+        # Encode vision inputs
+        encoder_hidden_states = self._encode_vision(pixel_values)
+        batch_size = pixel_values.shape[0] // 5
+        
+        # Start with decoder_start_token_id
+        decoder_input_ids = torch.full(
+            (batch_size, 1),
+            self._get_decoder_start_token_id(),
+            dtype=torch.long,
+            device=pixel_values.device
+        )
+        
+        generated_tokens = []
+        past_key_values = None
+        
+        for step in range(max_length):
+            with torch.no_grad():
+                # Get decoder outputs
+                decoder_outputs = self.seq2seq_decoder(
+                    input_ids=decoder_input_ids if past_key_values is None else decoder_input_ids[:, -1:],
+                    encoder_hidden_states=encoder_hidden_states,
+                    past_key_values=past_key_values,
+                    use_cache=True,
+                )
+                
+                past_key_values = decoder_outputs.past_key_values
+                
+                # Get logits and generate next token
+                hidden_states = decoder_outputs[0][:, -1:, :]
+                lm_logits = self.seq2seq_lm_head(hidden_states)
+                
+                # Apply temperature
+                if temperature != 1.0:
+                    lm_logits = lm_logits / temperature
+                
+                # Get next token
+                if do_sample:
+                    probs = torch.softmax(lm_logits[:, -1, :], dim=-1)
+                    next_token = torch.multinomial(probs, num_samples=1)
+                else:
+                    next_token = torch.argmax(lm_logits[:, -1, :], dim=-1, keepdim=True)
+                
+                # Append to generated tokens
+                generated_tokens.append(next_token)
+                decoder_input_ids = torch.cat([decoder_input_ids, next_token], dim=1)
+                
+                # Check for EOS
+                if (next_token == self.config.eos_token_id).all():
+                    break
+        
+        return decoder_input_ids
+    
+    def _get_decoder_start_token_id(self) -> int:
+        """Get decoder start token ID."""
+        return self.config.decoder_start_token_id or self.config.pad_token_id

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:aee68e205cbe66657917c8b719a66e0798425d2ba696ab9e23f81b6f8bbb7875
+size 758546423