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 from huggingface_hub import snapshot_download # import traceback login(os.environ.get("HUGGINGFACE_TOKEN")) os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:80" repo_id = "mistralai/Pixtral-12B-2409" max_tokens_per_img = 2048 max_img_per_msg = 2 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) model_name = "mistralai/Pixtral-12B-2409" llm = LLM( model=model_name, tokenizer_mode="mistral", max_num_batched_tokens=max_img_per_msg * max_tokens_per_img, dtype="float16" ) def clear_cuda_cache(): torch.cuda.empty_cache() 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 def infer(image_url, prompt, temperature, max_tokens, progress=gr.Progress(track_tqdm=True)): if llm is None: return "Error: LLM initialization failed. Please try again later." try: sampling_params = SamplingParams(max_tokens=max_tokens, temperature=temperature) 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) clear_cuda_cache() return outputs[0].outputs[0].text except Exception as e: clear_cuda_cache() return f"Error during inference: {e}" def compare_images(image1_url, image2_url, prompt, temperature, max_tokens, progress=gr.Progress(track_tqdm=True)): if llm is None: return "Error: LLM initialization failed. Please try again later." try: sampling_params = SamplingParams(max_tokens=max_tokens, temperature=temperature) 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) clear_cuda_cache() return outputs[0].outputs[0].text except Exception as e: clear_cuda_cache() return f"Error during image comparison: {e}" 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() clear_cuda_cache() return similarity except Exception as e: clear_cuda_cache() 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 - Adjust the temperature and max tokens - 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 - Adjust the temperature and max tokens - 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") with gr.Row(): temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, label="Temperature") max_tokens = gr.Number(value=4096, label="Max Tokens") generate_button = gr.Button("Generate") output = gr.Text(label="Generated Text") generate_button.click(infer, inputs=[image_url, prompt, temperature, max_tokens], 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") with gr.Row(): comparison_temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, label="Temperature") comparison_max_tokens = gr.Number(value=4096, label="Max Tokens") compare_button = gr.Button("Compare") comparison_output = gr.Text(label="Comparison Result") compare_button.click(compare_images, inputs=[image1_url, image2_url, comparison_prompt, comparison_temperature, comparison_max_tokens], 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()