app.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from safetensors import safe_open
import json
import gradio as gr
from PIL import Image
import numpy as np
from huggingface_hub import snapshot_download
from mistral_common.protocol.instruct.messages import UserMessage, TextChunk, ImageChunk
from mistral_common.protocol.instruct.request import ChatCompletionRequest
from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
import spaces
import math
from typing import List, Optional, Tuple
import gc
from contextlib import contextmanager
import os
from loadimg import load_img
import traceback

title = "# **WIP / DEMO** 🙋🏻‍♂️Welcome to Tonic's Pixtral Model Demo"
description = """
This demo showcases two capabilities of the Pixtral model:
1. Image-to-Text Generation
2. Image Similarity Comparison

### 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)

class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-5):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight

def precompute_freqs_cis_2d(dim: int, height: int, width: int, theta: float) -> torch.Tensor:
    freqs = 1.0 / (theta**(torch.arange(0, dim, 2).float() / dim))
    h = torch.arange(height)
    w = torch.arange(width)
    freqs_h = torch.outer(h, freqs[::2]).float()
    freqs_w = torch.outer(w, freqs[1::2]).float()
    freqs_2d = torch.cat([freqs_h[:, None, :].repeat(1, width, 1), freqs_w[None, :, :].repeat(height, 1, 1)], dim=-1)
    return torch.polar(torch.ones_like(freqs_2d), freqs_2d)

def apply_rotary_emb_vit(xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
    freqs_cis = freqs_cis.view(*freqs_cis.shape[:2], 1, freqs_cis.shape[-1])
    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
    return xq_out.type_as(xq), xk_out.type_as(xk)

class Attention(nn.Module):
    def __init__(self, args):
        super().__init__()
        self.n_heads = args['num_attention_heads']
        self.head_dim = args['hidden_size'] // args['num_attention_heads']
        self.wq = nn.Linear(args['hidden_size'], args['hidden_size'], bias=False)
        self.wk = nn.Linear(args['hidden_size'], args['hidden_size'], bias=False)
        self.wv = nn.Linear(args['hidden_size'], args['hidden_size'], bias=False)
        self.wo = nn.Linear(args['hidden_size'], args['hidden_size'], bias=False)

    def forward(self, x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
        batch, patches, _ = x.shape
        q, k, v = self.wq(x), self.wk(x), self.wv(x)
        q = q.reshape(batch, patches, self.n_heads, self.head_dim)
        k = k.reshape(batch, patches, self.n_heads, self.head_dim)
        v = v.reshape(batch, patches, self.n_heads, self.head_dim)
        q, k = apply_rotary_emb_vit(q, k, freqs_cis=freqs_cis)
        scores = torch.matmul(q, k.transpose(-1, -2)) / math.sqrt(self.head_dim)
        attn = F.softmax(scores, dim=-1)
        out = torch.matmul(attn, v)
        out = out.reshape(batch, patches, self.n_heads * self.head_dim)
        return self.wo(out)

class FeedForward(nn.Module):
    def __init__(self, args):
        super().__init__()
        self.w1 = nn.Linear(args['hidden_size'], args['intermediate_size'], bias=False)
        self.w2 = nn.Linear(args['intermediate_size'], args['hidden_size'], bias=False)
        self.w3 = nn.Linear(args['hidden_size'], args['intermediate_size'], bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

class TransformerBlock(nn.Module):
    def __init__(self, args):
        super().__init__()
        self.attention = Attention(args)
        self.feed_forward = FeedForward(args)
        self.attention_norm = RMSNorm(args['hidden_size'], eps=1e-5)
        self.ffn_norm = RMSNorm(args['hidden_size'], eps=1e-5)

    def forward(self, x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
        r = self.attention(self.attention_norm(x), freqs_cis=freqs_cis)
        h = x + r
        r = self.feed_forward(self.ffn_norm(h))
        out = h + r
        return out

class VisionTransformer(nn.Module):
    def __init__(self, args):
        super().__init__()
        self.args = args
        self.patch_conv = nn.Conv2d(args['num_channels'], args['hidden_size'], kernel_size=args['patch_size'], stride=args['patch_size'], bias=False)
        self.ln_pre = RMSNorm(args['hidden_size'], eps=1e-5)
        self.transformer = nn.ModuleList([TransformerBlock(args) for _ in range(args['num_hidden_layers'])])
        self.max_patches_per_side = args['image_size'] // args['patch_size']
        self._freqs_cis = None

    @property
    def freqs_cis(self) -> torch.Tensor:
        if self._freqs_cis is None:
            self._freqs_cis = precompute_freqs_cis_2d(
                dim=self.args['hidden_size'] // self.args['num_attention_heads'],
                height=self.max_patches_per_side,
                width=self.max_patches_per_side,
                theta=self.args['rope_theta'],
            )
        return self._freqs_cis.to(self.patch_conv.weight.device)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.patch_conv(x)
        x = x.flatten(2).transpose(1, 2)
        x = self.ln_pre(x)
        freqs_cis = self.freqs_cis
        for layer in self.transformer:
            x = layer(x, freqs_cis=freqs_cis)
        return x

class VisionLanguageAdapter(nn.Module):
    def __init__(self, args, dim: int):
        super().__init__()
        self.w_in = nn.Linear(args['hidden_size'], dim, bias=True)
        self.gelu = nn.GELU()
        self.w_out = nn.Linear(dim, dim, bias=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.w_out(self.gelu(self.w_in(x)))

class PixtralModel(nn.Module):
    def __init__(self, params):
        super().__init__()
        self.vision_encoder = VisionTransformer(params['vision_encoder'])
        self.vision_language_adapter = VisionLanguageAdapter(params['vision_encoder'], params['dim'])
        self.language_model = nn.TransformerDecoder(
            nn.TransformerDecoderLayer(d_model=params['dim'], nhead=params['n_heads'], dim_feedforward=params['hidden_dim']),
            num_layers=params['n_layers']
        )
        self.lm_head = nn.Linear(params['dim'], params['vocab_size'], bias=False)

    def forward(self, image, input_ids=None):
        vision_output = self.vision_encoder(image)
        vision_output = self.vision_language_adapter(vision_output)
        
        if input_ids is not None:
            tgt = self.lm_head.weight[input_ids].transpose(0, 1)
            output = self.language_model(tgt, vision_output)
            logits = self.lm_head(output)
            return logits
        else:
            return vision_output


@contextmanager
def gpu_memory_manager():
    try:
        torch.cuda.empty_cache()
        yield
    finally:
        torch.cuda.empty_cache()
        gc.collect()

def load_model_with_fallback(params, model_path):
    try:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        model = PixtralModel(params)
        with safe_open(f'{model_path}/consolidated.safetensors', framework="pt", device="cpu") as f:
            for name, param in model.named_parameters():
                if name in f.keys():
                    param.data = f.get_tensor(name)
        model.eval()
        model.to(device)
        return model, device
    except RuntimeError as e:
        print(f"Error loading model on GPU: {str(e)}")
        print("Falling back to CPU...")
        model = PixtralModel(params)
        with safe_open(f'{model_path}/consolidated.safetensors', framework="pt", device="cpu") as f:
            for name, param in model.named_parameters():
                if name in f.keys():
                    param.data = f.get_tensor(name)
        model.eval()
        return model, torch.device("cpu")

model, device = load_model_with_fallback(params, model_path)
tokenizer = MistralTokenizer.from_model("pixtral")

def preprocess_image(image):
    if image is None:
        raise ValueError("No image provided")
    
    pil_image = load_img(image, output_type="pil", input_type="auto")
    
    pil_image = pil_image.convert('RGB')
    pil_image = pil_image.resize((params['vision_encoder']['image_size'], params['vision_encoder']['image_size']))
    image_tensor = torch.tensor(np.array(pil_image)).permute(2, 0, 1).unsqueeze(0).float() / 255.0
    return image_tensor

@contextmanager
def gpu_memory_manager():
    try:
        torch.cuda.empty_cache()
        yield
    finally:
        torch.cuda.empty_cache()
        gc.collect()

def cuda_error_handler(func):
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except RuntimeError as e:
            if "CUDA" in str(e):
                print(f"CUDA error occurred: {str(e)}")
                print("Attempting to recover...")
                torch.cuda.empty_cache()
                gc.collect()
                try:
                    return func(*args, **kwargs)
                except Exception as e2:
                    print(f"Recovery failed. Error: {str(e2)}")
                    return f"An error occurred: {str(e2)}", 0, 0
            else:
                raise
        except Exception as e:
            print(f"An unexpected error occurred: {str(e)}")
            traceback.print_exc()
            return f"An unexpected error occurred: {str(e)}", 0, 0
    return wrapper

@spaces.GPU(duration=120)
@cuda_error_handler
def generate_text(image, prompt, max_tokens):
    try:
        with gpu_memory_manager():
            image_pil = load_img(image, output_type="pil", input_type="auto")
            image_tensor = preprocess_image(image_pil).to(device)

            tokenized = tokenizer.encode_chat_completion(
                ChatCompletionRequest(
                    messages=[UserMessage(content=[TextChunk(text=prompt), ImageChunk(image=image)])],
                    model="pixtral",
                )
            )
            input_ids = torch.tensor(tokenized.tokens).unsqueeze(0).to(device)

            generated_ids = input_ids.clone()
            for _ in range(max_tokens):
                with torch.no_grad():
                    logits = model(image_tensor, generated_ids)
                next_token_logits = logits[0, -1, :]
                next_token = torch.argmax(next_token_logits, dim=-1)
                generated_ids = torch.cat([generated_ids, next_token.unsqueeze(0).unsqueeze(0)], dim=-1)
                if next_token.item() == tokenizer.eos_token_id:
                    break

            generated_text = tokenizer.decode(generated_ids[0].tolist())
            
        torch.cuda.empty_cache()
        
        return generated_text, len(generated_ids[0]), 1
    except Exception as e:
        print(f"Error in generate_text: {str(e)}")
        traceback.print_exc()
        return f"Error: {str(e)}", 0, 0

@spaces.GPU(duration=60)
@cuda_error_handler
def calculate_similarity(image1, image2):
    try:
        with gpu_memory_manager():
            pil_image1 = load_img(image1, output_type="pil", input_type="auto")
            pil_image2 = load_img(image2, output_type="pil", input_type="auto")
            tensor1 = preprocess_image(pil_image1).to(device)
            tensor2 = preprocess_image(pil_image2).to(device)

            with torch.no_grad():
                embedding1 = model(tensor1).mean(dim=1)
                embedding2 = model(tensor2).mean(dim=1)

            similarity = F.cosine_similarity(embedding1, embedding2).item()
        
        torch.cuda.empty_cache()
        
        return similarity
    except Exception as e:
        print(f"Error in calculate_similarity: {str(e)}")
        traceback.print_exc()
        return f"Error: {str(e)}"
    
# @spaces.GPU()
# @cuda_error_handler
# def calculate_similarity(image1, image2):
#     try:
#         with gpu_memory_manager():
#             device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
#             # Use load_img for both images
#             pil_image1 = load_img(image1, output_type="pil", input_type="auto")
#             pil_image2 = load_img(image2, output_type="pil", input_type="auto")
#             tensor1 = preprocess_image(pil_image1).to(device)
#             tensor2 = preprocess_image(pil_image2).to(device)
#             model.to(device)

#             with torch.no_grad():
#                 embedding1 = model(tensor1).mean(dim=1)
#                 embedding2 = model(tensor2).mean(dim=1)

#             similarity = F.cosine_similarity(embedding1, embedding2).item()
        
#         # # Move model back to CPU and clear CUDA memory
#         # model.to("cpu")
#         torch.cuda.empty_cache()
        
#         return similarity
#     except Exception as e:
#         print(f"Error in calculate_similarity: {str(e)}")
#         traceback.print_exc()
#         return f"Error: {str(e)}"

with gr.Blocks() as demo:
    gr.Markdown(title)
    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']}")
    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(description)
    
    with gr.Tabs():
        with gr.TabItem("Image-to-Text Generation"):
            with gr.Row():
                with gr.Column():
                    input_image = gr.Image(type="pil", label="Input Image")
                    input_prompt = gr.Textbox(label="Prompt")
                    max_tokens_slider = gr.Slider(minimum=10, maximum=500, value=100, step=10, label="Max Tokens")
                    submit_btn = gr.Button("Generate Text")
                
                with gr.Column():
                    output_text = gr.Textbox(label="Generated Text")
                    token_count = gr.Number(label="Number of Tokens")
                    image_count = gr.Number(label="Number of Images Processed")
            
            submit_btn.click(
                fn=generate_text,
                inputs=[input_image, input_prompt, max_tokens_slider],
                outputs=[output_text, token_count, image_count]
            )
        
        with gr.TabItem("Image Similarity Comparison"):
            with gr.Row():
                image1_input = gr.Image(type="pil", label="Image 1")
                image2_input = gr.Image(type="pil", label="Image 2")
            
            similarity_btn = gr.Button("📸🌬️Calculate Similarity")
            similarity_output = gr.Number(label="Similarity Score (0.0 to 1.0)")
            
            similarity_btn.click(
                fn=calculate_similarity,
                inputs=[image1_input, image2_input],
                outputs=[similarity_output]
            )

if __name__ == "__main__":
    try:
        demo.launch()
    except Exception as e:
        print(f"An error occurred while launching the demo: {str(e)}")
        traceback.print_exc()