Create app.py

app.py

@@ -0,0 +1,129 @@
+import gradio as gr
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel, PeftConfig
+from PIL import Image
+import requests
+from io import BytesIO
+import torchvision.datasets as datasets
+import numpy as np
+
+# Load SigLIP for image embeddings
+from model.siglip import SigLIPModel
+
+def get_cifar_examples():
+    # Load CIFAR10 test set
+    cifar10_test = datasets.CIFAR10(root='./data', train=False, download=True)
+    
+    # CIFAR10 classes
+    classes = ['airplane', 'automobile', 'bird', 'cat', 'deer', 
+              'dog', 'frog', 'horse', 'ship', 'truck']
+    
+    # Get one example from each class
+    examples = []
+    used_classes = set()
+    
+    for idx in range(len(cifar10_test)):
+        img, label = cifar10_test[idx]
+        if classes[label] not in used_classes:
+            # Save the image temporarily
+            img_path = f"examples/{classes[label]}_example.jpg"
+            img.save(img_path)
+            examples.append(img_path)
+            used_classes.add(classes[label])
+            
+        if len(used_classes) == 10:  # We have one example from each class
+            break
+    
+    return examples
+
+def load_models():
+    # Load SigLIP model
+    siglip = SigLIPModel()
+    
+    # Load base Phi model
+    base_model = AutoModelForCausalLM.from_pretrained(
+        "microsoft/Phi-3-mini-4k-instruct",
+        trust_remote_code=True,
+        device_map="auto",
+        torch_dtype=torch.float32
+    )
+    
+    # Load our fine-tuned LoRA adapter
+    model = PeftModel.from_pretrained(
+        base_model,
+        "jatingocodeo/phi-vlm",  # Your uploaded model
+        device_map="auto"
+    )
+    
+    tokenizer = AutoTokenizer.from_pretrained("jatingocodeo/phi-vlm")
+    
+    return siglip, model, tokenizer
+
+def generate_description(image, siglip, model, tokenizer):
+    # Convert image to RGB if needed
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+    
+    # Resize image to match SigLIP's expected size
+    image = image.resize((32, 32))
+    
+    # Get image embedding from SigLIP
+    image_embedding = siglip.encode_image(image)
+    
+    # Prepare prompt
+    prompt = """Below is an image. Please describe it in detail.
+
+Image: <image>
+Description: """
+    
+    # Tokenize input
+    inputs = tokenizer(
+        prompt,
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+        max_length=128
+    ).to(model.device)
+    
+    # Generate description
+    with torch.no_grad():
+        outputs = model(
+            **inputs,
+            image_embeddings=image_embedding.unsqueeze(0),
+            max_new_tokens=100,
+            temperature=0.7,
+            do_sample=True,
+            top_p=0.9
+        )
+    
+    # Decode and return the generated text
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return generated_text.split("Description: ")[-1].strip()
+
+# Load models
+print("Loading models...")
+siglip, model, tokenizer = load_models()
+
+# Create Gradio interface
+def process_image(image):
+    description = generate_description(image, siglip, model, tokenizer)
+    return description
+
+# Get CIFAR10 examples
+examples = get_cifar_examples()
+
+# Define interface
+iface = gr.Interface(
+    fn=process_image,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Textbox(label="Generated Description"),
+    title="Image Description Generator",
+    description="""Upload an image and get a detailed description generated by our fine-tuned VLM model.
+                  Below are sample images from CIFAR10 dataset that you can try.""",
+    examples=[[ex] for ex in examples]  # Format examples for Gradio
+)
+
+# Launch the interface
+if __name__ == "__main__":
+    iface.launch()