# app.py import spaces import os import gradio as gr import torch from PIL import Image from transformers import AutoTokenizer, AutoModelForCausalLM,AutoConfig, BitsAndBytesConfig import timm from torchvision import transforms #from llama_cpp import Llama from peft import PeftModel, prepare_model_for_kbit_training, LoraConfig, get_peft_model, TaskType import traceback import warnings warnings.filterwarnings('ignore') # 1. Model Definitions (Same as in training script) class SigLIPImageEncoder(torch.nn.Module): def __init__(self, model_name='resnet50', embed_dim=512, pretrained_path=None): super().__init__() self.model = timm.create_model(model_name, pretrained=False, num_classes=0, global_pool='avg') # pretrained=False self.embed_dim = embed_dim self.projection = torch.nn.Linear(self.model.num_features, embed_dim) if pretrained_path: #self.load_state_dict(torch.load(pretrained_path, map_location=torch.device('cpu'))) # Load to CPU first self.load_state_dict(torch.load(pretrained_path)) print(f"Loaded SigLIP image encoder from {pretrained_path}") else: print("Initialized SigLIP image encoder without pretrained weights.") def forward(self, image): features = self.model(image) embedding = self.projection(features) return embedding class Phi3WithImage(torch.nn.Module): 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, torch_dtype=torch.float32, device_map="auto", 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 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] # 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 at the start inputs_embeds = torch.cat([projected_image_embeddings, token_embeddings], dim=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) # 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) 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" image_model_name = 'resnet50' image_embed_dim = 512 siglip_pretrained_path = "image_encoder.pth" #device = torch.device("cpu") # Force CPU device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Using device: {device}") # Load Tokenizer 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)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) ]) # 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, bnb_config=bnb_config).to(device) model = Phi3WithImage(phi3_model_name, image_encoder, image_embed_dim, bnb_config=bnb_config).to(device) #model.phi3 = prepare_model_for_kbit_training(model.phi3) # Disable cache manually (important!) if hasattr(model.phi3.config, 'use_cache'): model.phi3.config.use_cache = False # Load LoRA model model.phi3 = PeftModel.from_pretrained(model.phi3, lora_model_path, device_map="auto", offload_dir='./offload') model.phi3 = model.phi3.merge_and_unload() #model.phi3.save_pretrained("merged_model_fp16") #model.phi3 = PeftModel.from_pretrained(model.phi3, lora_model_path, device_map="auto") #model.phi3 = AutoModelForCausalLM.from_pretrained( # "merged_model_fp16", # quantization_config=bnb_config, # torch_dtype=torch.bfloat16, # device_map="auto" #) model.eval() # Set to evaluation mode # 3. Inference Function @spaces.GPU # app.py # ... existing code ... # 3. Inference Function def predict(image_input, question): """ Takes an image and a question as input and returns an answer. """ if image_input is None or question is None or question == "": return "Please provide both an image and a question." try: image = Image.fromarray(image_input).convert("RGB") image = image_transform(image).unsqueeze(0).to(device) prompt = f"Question: {question}\nAnswer:" encoded = text_tokenizer(prompt, return_tensors="pt").to(device) with torch.no_grad(): # Get image embeddings image_embeddings = model.image_encoder(image) projected_image_embeddings = model.image_projection(image_embeddings) # Reshape image embeddings to (batch_size, 1, phi3_embed_dim) projected_image_embeddings = projected_image_embeddings.unsqueeze(1) # Concatenate along the sequence dimension (dim=1) extended_attention_mask = torch.cat([torch.ones(projected_image_embeddings.shape[:2], device=encoded["attention_mask"].device), encoded["attention_mask"]], dim=1) extended_input_ids = torch.cat([torch.zeros(projected_image_embeddings.shape[:2], dtype=torch.long, device=encoded["input_ids"].device), encoded["input_ids"]], dim=1) # Generate answer generated_tokens = model.phi3.generate( input_ids=extended_input_ids, attention_mask=extended_attention_mask, max_length=128, pad_token_id=text_tokenizer.eos_token_id, ) answer = text_tokenizer.decode(generated_tokens[0], skip_special_tokens=True) answer = answer.replace(prompt, "").strip() # Remove prompt from answer return answer except Exception as e: #return f"An error occurred: {str(e)}" return f"An error occurred: {traceback.format_exc()}" # 5. Launch the App if __name__ == "__main__": iface = gr.Interface( fn=predict, inputs=[ gr.Image(label="Upload an Image"), gr.Textbox(label="Ask a Question about the Image", placeholder="What is in the image?") ], outputs=gr.Textbox(label="Answer"), title="Image Question Answering with Phi-3 and SigLIP", description="Ask questions about an image and get answers powered by Phi-3 and SigLIP.", examples=[ ["cat_0006.png", "Create a interesting story about this image?"], ["bird_0004.png", "Can you describe this image?"], ["truck_0003.png", "Elaborate the setting of the image"], ["ship_0007.png", "Explain the purpose of image"] ] ) iface.launch()