Update app.py

app.py

@@ -4,7 +4,7 @@ import os
 import gradio as gr
 import torch
 from PIL import Image
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM,AutoConfig, BitsAndBytesConfig
 import timm
 from torchvision import transforms
 #from llama_cpp import Llama
@@ -31,42 +31,53 @@ class SigLIPImageEncoder(torch.nn.Module):
         embedding = self.projection(features)
         return embedding
 
+
 class Phi3WithImage(torch.nn.Module):
-    def __init__(self, phi3_model_name, lora_model_path, image_encoder, image_embed_dim=512):
+    
+    def __init__(self, phi3_model_name, image_encoder, image_embed_dim=512, image_token_id=1, bnb_config=None): 
         super().__init__()
         self.phi3 = AutoModelForCausalLM.from_pretrained(
             phi3_model_name,
             torch_dtype=torch.bfloat16,
             device_map="auto",
-            trust_remote_code=True  # Important for some Phi-3 variants
+            trust_remote_code=True,
+            quantization_config=bnb_config
         )
-
         self.image_encoder = image_encoder
         self.image_embed_dim = image_embed_dim
         self.phi3_embed_dim = self.phi3.config.hidden_size
-
-        # Project image embeddings to Phi-3's embedding space
         self.image_projection = torch.nn.Linear(image_embed_dim, self.phi3_embed_dim)
+        self.image_token_id = image_token_id
+
+    def forward(self, image, question_input_ids, question_attention_mask, answer_input_ids, answer_attention_mask):
+        batch_size = question_input_ids.size(0)
+
+        # Encode image and project to text embedding space
+        image_embeddings = self.image_encoder(image)  # Shape: [batch_size, image_embed_dim]
+        projected_image_embeddings = self.image_projection(image_embeddings).unsqueeze(1)  # Shape: [batch_size, 1, hidden_dim]
 
-    def forward(self, image, question_input_ids, question_attention_mask):
-        image_embeddings = self.image_encoder(image)
-        projected_image_embeddings = self.image_projection(image_embeddings)
+        # Get the token embeddings for question + answer from the model’s embedding layer
+        token_embeddings = self.phi3.get_input_embeddings()(
+            torch.cat([question_input_ids, answer_input_ids], dim=1)
+        )  # Shape: [batch_size, seq_len, hidden_dim]
 
-        # Concatenate image embeddings to the input sequence
-        #  This is a simplified approach.  More sophisticated methods exist.
-        #  Assumes image embeddings are prepended to the sequence.
+        # Concatenate image embeddings at the start
+        inputs_embeds = torch.cat([projected_image_embeddings, token_embeddings], dim=1)
 
-        #  Reshape image embeddings to (batch_size, 1, phi3_embed_dim)
-        projected_image_embeddings = projected_image_embeddings.unsqueeze(1)
+        # Create combined attention mask
+        image_attention_mask = torch.ones((batch_size, 1), device=question_attention_mask.device)
+        full_attention_mask = torch.cat([image_attention_mask, question_attention_mask, answer_attention_mask], dim=1)
 
-        #  Concatenate along the sequence dimension (dim=1)
-        extended_attention_mask = torch.cat([torch.ones(projected_image_embeddings.shape[:2], device=question_attention_mask.device), question_attention_mask], dim=1)
-        extended_input_ids = torch.cat([torch.zeros(projected_image_embeddings.shape[:2], dtype=torch.long, device=question_input_ids.device), question_input_ids], dim=1)
+        # Prepare labels: mask image + question parts so only answers are supervised
+        #labels = torch.cat([torch.full((batch_size, 1 + question_input_ids.size(1)), -100, device=question_input_ids.device), answer_input_ids], dim=1)
 
-        #  Pass the concatenated input to Phi-3
-        outputs = self.phi3(input_ids=extended_input_ids, attention_mask=extended_attention_mask) # No labels during inference
+        outputs = self.phi3(
+            inputs_embeds=inputs_embeds,
+            attention_mask=full_attention_mask
+        )
         return outputs.logits
 
+
 # 2. Load Models and Tokenizer
 phi3_model_name = "microsoft/Phi-3-mini-4k-instruct"  # Or your specific Phi-3 variant
 lora_model_path = "./qlora-phi3-model-new"
@@ -81,6 +92,14 @@ print(f"Using device: {device}")
 text_tokenizer = AutoTokenizer.from_pretrained(phi3_model_name, trust_remote_code=True)
 text_tokenizer.pad_token = text_tokenizer.eos_token # Important for training
 
+# 7. BitsAndBytesConfig for QLoRA
+bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_use_double_quant=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.bfloat16
+        )
+
 # Image Transformations
 image_transform = transforms.Compose([
     transforms.Resize((224, 224)),
@@ -90,7 +109,8 @@ image_transform = transforms.Compose([
 
 # Load Models
 image_encoder = SigLIPImageEncoder(model_name=image_model_name, embed_dim=image_embed_dim, pretrained_path=siglip_pretrained_path).to(device)
-model = Phi3WithImage(phi3_model_name, lora_model_path, image_encoder, image_embed_dim).to(device)
+#model = Phi3WithImage(phi3_model_name, lora_model_path, image_encoder, image_embed_dim, bnb_config=bnb_config).to(device)
+model = Phi3WithImage(phi3_model_name,  image_encoder, image_embed_dim, bnb_config=bnb_config).to(device)
 
 # Load LoRA model
 model.phi3 = PeftModel.from_pretrained(model.phi3, lora_model_path, device_map="auto", offload_dir='./offload')