Create app.py

app.py

@@ -0,0 +1,116 @@
+import gradio as gr
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from transformers import AutoModel, AutoTokenizer, AutoImageProcessor
+
+# Check if flash_attn is available
+def is_flash_attn_available():
+    try:
+        import flash_attn
+        return True
+    except ImportError:
+        return False
+
+# Load model and tokenizer
+@torch.inference_mode()
+def load_model():
+    use_optimized = torch.cuda.is_available() and is_flash_attn_available()
+    
+    model = AutoModel.from_pretrained(
+        "visheratin/mexma-siglip2", 
+        torch_dtype=torch.bfloat16, 
+        trust_remote_code=True, 
+        optimized=True if use_optimized else False,
+    )
+    if torch.cuda.is_available():
+        model = model.to("cuda")
+    tokenizer = AutoTokenizer.from_pretrained("visheratin/mexma-siglip2")
+    processor = AutoImageProcessor.from_pretrained("visheratin/mexma-siglip2")
+    return model, tokenizer, processor
+
+model, tokenizer, processor = load_model()
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+def classify_image(image, text_queries):
+    if image is None or not text_queries.strip():
+        return None
+    
+    # Process image
+    processed_image = processor(images=image, return_tensors="pt")["pixel_values"]
+    processed_image = processed_image.to(torch.bfloat16)
+    if torch.cuda.is_available():
+        processed_image = processed_image.to("cuda")
+    
+    # Process text queries
+    queries = [q.strip() for q in text_queries.split("\n") if q.strip()]
+    if not queries:
+        return None
+    
+    text_inputs = tokenizer(queries, return_tensors="pt", padding=True)
+    if torch.cuda.is_available():
+        text_inputs = text_inputs.to("cuda")
+    
+    # Get predictions
+    with torch.inference_mode():
+        image_logits, _ = model.get_logits(
+            text_inputs["input_ids"], 
+            text_inputs["attention_mask"], 
+            processed_image
+        )
+        probs = F.softmax(image_logits, dim=-1)[0].cpu().tolist()
+    
+    # Format results
+    results = {queries[i]: f"{probs[i]:.4f}" for i in range(len(queries))}
+    return results
+
+# Create Gradio interface
+with gr.Blocks(title="Mexma-SigLIP2 Zero-Shot Classification") as demo:
+    gr.Markdown("# Mexma-SigLIP2 Zero-Shot Classification Demo")
+    gr.Markdown("""
+    This demo showcases the zero-shot classification capabilities of the Mexma-SigLIP2 model.
+    
+    ### Instructions:
+    1. Upload or select an image
+    2. Enter text queries (one per line) to classify the image
+    3. Click 'Submit' to see the classification probabilities
+    
+    The model supports multilingual queries (English, Russian, Hindi, etc.)
+    """)
+    
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.Image(type="pil", label="Upload Image")
+            text_input = gr.Textbox(
+                placeholder="Enter text queries (one per line)\nExample:\na cat\na dog\nEiffel Tower", 
+                label="Text Queries",
+                lines=5
+            )
+            submit_btn = gr.Button("Submit", variant="primary")
+        
+        with gr.Column():
+            output = gr.Label(label="Classification Results")
+    
+    submit_btn.click(
+        fn=classify_image,
+        inputs=[image_input, text_input],
+        outputs=output
+    )
+    
+    gr.Examples(
+        [
+            [
+                "https://static.independent.co.uk/s3fs-public/thumbnails/image/2014/03/25/12/eiffel.jpg",
+                "Eiffel Tower\nStatue of Liberty\nTaj Mahal\nкошка\nएफिल टॉवर"
+            ],
+            [
+                "https://upload.wikimedia.org/wikipedia/commons/thumb/3/3a/Cat03.jpg/1200px-Cat03.jpg",
+                "a cat\na dog\na bird\nкошка\nсобака"
+            ]
+        ],
+        inputs=[image_input, text_input]
+    )
+
+# Launch the demo
+if __name__ == "__main__":
+    demo.launch()