app.py

import os
import gradio as gr
from vllm import LLM, SamplingParams
from PIL import Image
from io import BytesIO
import base64
import requests
from huggingface_hub import login
import torch
import torch.nn.functional as F
import spaces
import json
import gradio as gr
from huggingface_hub import snapshot_download
import os
# from loadimg import load_img
import traceback

login(os.environ.get("HUGGINGFACE_TOKEN"))

repo_id = "mistralai/Pixtral-12B-2409"
sampling_params = SamplingParams(max_tokens=8192, temperature=0.7)
max_tokens_per_img = 4096
max_img_per_msg = 5


title = "# **WIP / DEMO** 🙋🏻‍♂️Welcome to Tonic's Pixtral Model Demo"
description = """
### Join us : 
🌟TeamTonic🌟 is always making cool demos! Join our active builder's 🛠️community 👻 [![Join us on Discord](https://img.shields.io/discord/1109943800132010065?label=Discord&logo=discord&style=flat-square)](https://discord.gg/qdfnvSPcqP) On 🤗Huggingface:[MultiTransformer](https://huggingface.co./MultiTransformer) On 🌐Github: [Tonic-AI](https://github.com/tonic-ai) & contribute to🌟 [Build Tonic](https://git.tonic-ai.com/contribute)🤗Big thanks to Yuvi Sharma and all the folks at huggingface for the community grant 🤗
"""
HUGGINGFACE_TOKEN = os.environ.get("HUGGINGFACE_TOKEN")
model_path = snapshot_download(repo_id="mistralai/Pixtral-12B-2409", token=HUGGINGFACE_TOKEN)

with open(f'{model_path}/params.json', 'r') as f:
    params = json.load(f)

with open(f'{model_path}/tekken.json', 'r') as f:
    tokenizer_config = json.load(f)

@spaces.GPU()
def initialize_llm():
    try:
        llm = LLM(
            model=repo_id,
            tokenizer_mode="mistral",
            max_model_len=65536,
            max_num_batched_tokens=max_img_per_msg * max_tokens_per_img,
            limit_mm_per_prompt={"image": max_img_per_msg}
        )
        return llm
    except Exception as e:
        print("LLM initialization failed:", e)
        return None

sampling_params = SamplingParams(max_tokens=8192)
llm = initialize_llm()

def encode_image(image: Image.Image, image_format="PNG") -> str:
    im_file = BytesIO()
    image.save(im_file, format=image_format)
    im_bytes = im_file.getvalue()
    im_64 = base64.b64encode(im_bytes).decode("utf-8")
    return im_64

@spaces.GPU()
def infer(image_url, prompt, progress=gr.Progress(track_tqdm=True)):
    if llm is None:
        return "Error: LLM initialization failed. Please try again later."
    
    try:
        image = Image.open(BytesIO(requests.get(image_url).content))
        image = image.resize((3844, 2408))
        new_image_url = f"data:image/png;base64,{encode_image(image, image_format='PNG')}"

        messages = [
            {
                "role": "user",
                "content": [{"type": "text", "text": prompt}, {"type": "image_url", "image_url": {"url": new_image_url}}]
            },
        ]

        outputs = llm.chat(messages, sampling_params=sampling_params)

        return outputs[0].outputs[0].text
    except Exception as e:
        return f"Error during inference: {e}"

@spaces.GPU()
def compare_images(image1_url, image2_url, prompt, progress=gr.Progress(track_tqdm=True)):
    if llm is None:
        return "Error: LLM initialization failed. Please try again later."
    
    try:
        image1 = Image.open(BytesIO(requests.get(image1_url).content))
        image2 = Image.open(BytesIO(requests.get(image2_url).content))
        image1 = image1.resize((3844, 2408))
        image2 = image2.resize((3844, 2408))
        new_image1_url = f"data:image/png;base64,{encode_image(image1, image_format='PNG')}"
        new_image2_url = f"data:image/png;base64,{encode_image(image2, image_format='PNG')}"

        messages = [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {"type": "image_url", "image_url": {"url": new_image1_url}},
                    {"type": "image_url", "image_url": {"url": new_image2_url}}
                ]
            },
        ]

        outputs = llm.chat(messages, sampling_params=sampling_params)

        return outputs[0].outputs[0].text
    except Exception as e:
        return f"Error during image comparison: {e}"

@spaces.GPU()
def calculate_image_similarity(image1_url, image2_url):
    if llm is None:
        return "Error: LLM initialization failed. Please try again later."
    
    try:
        image1 = Image.open(BytesIO(requests.get(image1_url).content)).convert('RGB')
        image2 = Image.open(BytesIO(requests.get(image2_url).content)).convert('RGB')
        image1 = image1.resize((224, 224))  # Resize to match model input size
        image2 = image2.resize((224, 224))

        image1_tensor = torch.tensor(list(image1.getdata())).view(1, 3, 224, 224).float() / 255.0
        image2_tensor = torch.tensor(list(image2.getdata())).view(1, 3, 224, 224).float() / 255.0

        with torch.no_grad():
            embedding1 = llm.model.vision_encoder([image1_tensor])
            embedding2 = llm.model.vision_encoder([image2_tensor])

        similarity = F.cosine_similarity(embedding1.mean(dim=0), embedding2.mean(dim=0), dim=0).item()

        return similarity
    except Exception as e:
        return f"Error during image similarity calculation: {e}"

with gr.Blocks() as demo:
    gr.Markdown(title)
    gr.Markdown("## How it works")
    gr.Markdown("1. The image is processed by a Vision Encoder using 2D ROPE (Rotary Position Embedding).")
    gr.Markdown("2. The encoder uses SiLU activation in its feed-forward layers.")
    gr.Markdown("3. The encoded image is used for text generation or similarity comparison.")
    gr.Markdown(
        """
        ## How to use
        1. For Image-to-Text Generation:
           - Enter the URL of an image
           - Provide a prompt describing what you want to know about the image
           - Click "Generate" to get the model's response
        2. For Image Comparison:
           - Enter URLs for two images you want to compare
           - Provide a prompt asking about the comparison
           - Click "Compare" to get the model's analysis
        3. For Image Similarity:
           - Enter URLs for two images you want to compare
           - Click "Calculate Similarity" to get a similarity score between 0 and 1
        """
    )
    gr.Markdown(description)
    with gr.Tabs():
        with gr.TabItem("Image-to-Text Generation"):
            with gr.Row():
                image_url = gr.Text(label="Image URL")
                prompt = gr.Text(label="Prompt")
            generate_button = gr.Button("Generate")
            output = gr.Text(label="Generated Text")
            
            generate_button.click(infer, inputs=[image_url, prompt], outputs=output)
        
        with gr.TabItem("Image Comparison"):
            with gr.Row():
                image1_url = gr.Text(label="Image 1 URL")
                image2_url = gr.Text(label="Image 2 URL")
            comparison_prompt = gr.Text(label="Comparison Prompt")
            compare_button = gr.Button("Compare")
            comparison_output = gr.Text(label="Comparison Result")
            
            compare_button.click(compare_images, inputs=[image1_url, image2_url, comparison_prompt], outputs=comparison_output)

        with gr.TabItem("Image Similarity"):
            with gr.Row():
                sim_image1_url = gr.Text(label="Image 1 URL")
                sim_image2_url = gr.Text(label="Image 2 URL")
            similarity_button = gr.Button("Calculate Similarity")
            similarity_output = gr.Number(label="Similarity Score")

            similarity_button.click(calculate_image_similarity, inputs=[sim_image1_url, sim_image2_url], outputs=similarity_output)
    gr.Markdown("## Model Details")
    gr.Markdown(f"- Model Dimension: {params['dim']}")
    gr.Markdown(f"- Number of Layers: {params['n_layers']}")
    gr.Markdown(f"- Number of Attention Heads: {params['n_heads']}")
    gr.Markdown(f"- Vision Encoder Hidden Size: {params['vision_encoder']['hidden_size']}")
    gr.Markdown(f"- Number of Vision Encoder Layers: {params['vision_encoder']['num_hidden_layers']}")
    gr.Markdown(f"- Number of Vision Encoder Attention Heads: {params['vision_encoder']['num_attention_heads']}")
    gr.Markdown(f"- Image Size: {params['vision_encoder']['image_size']}x{params['vision_encoder']['image_size']}")
    gr.Markdown(f"- Patch Size: {params['vision_encoder']['patch_size']}x{params['vision_encoder']['patch_size']}")

if __name__ == "__main__":
    demo.launch()