add reference code from vllm

.gitignore

@@ -0,0 +1 @@
+notes.py

app.py

@@ -14,7 +14,7 @@ import spaces
 import math
 from typing import List, Optional, Tuple
 
-title = "# 🙋🏻‍♂️Welcome to Tonic's Pixtral Model Demo"
+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
@@ -25,6 +25,7 @@ This demo showcases two capabilities of the Pixtral model:
 """
 
 model_path = snapshot_download(repo_id="mistralai/Pixtral-12B-2409")
+
 with open(f'{model_path}/params.json', 'r') as f:
     params = json.load(f)
 
@@ -40,32 +41,16 @@ class RMSNorm(nn.Module):
     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:
+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, device=freqs.device)
-    w = torch.arange(width, device=freqs.device)
-
+    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,
-    )
+    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]:
+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])
@@ -78,7 +63,6 @@ class Attention(nn.Module):
         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)
@@ -86,14 +70,11 @@ class Attention(nn.Module):
 
     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)
@@ -119,9 +100,9 @@ class TransformerBlock(nn.Module):
         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.forward(self.attention_norm(x), freqs_cis=freqs_cis)
+        r = self.attention(self.attention_norm(x), freqs_cis=freqs_cis)
         h = x + r
-        r = self.feed_forward.forward(self.ffn_norm(h))
+        r = self.feed_forward(self.ffn_norm(h))
         out = h + r
         return out
 
@@ -129,16 +110,9 @@ class VisionTransformer(nn.Module):
     def __init__(self, args):
         super().__init__()
         self.args = args
-        self.patch_conv = nn.Conv2d(
-            in_channels=args['num_channels'],
-            out_channels=args['hidden_size'],
-            kernel_size=args['patch_size'],
-            stride=args['patch_size'],
-            bias=False,
-        )
+        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
 
@@ -157,11 +131,9 @@ class VisionTransformer(nn.Module):
         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):
@@ -180,9 +152,7 @@ class PixtralModel(nn.Module):
         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']),
+            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)
@@ -201,12 +171,10 @@ class PixtralModel(nn.Module):
 
 def load_model(params, model_path):
     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
 
@@ -224,53 +192,45 @@ def preprocess_image(image):
 @spaces.GPU(duration=120)
 def generate_text(image, prompt, max_tokens):
     try:
-        image_tensor = preprocess_image(image).cuda()
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        image_tensor = preprocess_image(image).to(device)
+        model.to(device)
 
         tokenized = tokenizer.encode_chat_completion(
             ChatCompletionRequest(
-                messages=[
-                    UserMessage(
-                        content=[
-                            TextChunk(text=prompt),
-                            ImageChunk(image=image),
-                        ]
-                    )
-                ],
+                messages=[UserMessage(content=[TextChunk(text=prompt), ImageChunk(image=image)])],
                 model="pixtral",
             )
         )
-        input_ids = torch.tensor(tokenized.tokens).unsqueeze(0).cuda()
-
-        with torch.no_grad():
-            model.cuda()
-            for _ in range(max_tokens):
-                logits = model(image_tensor, input_ids)
-                next_token_logits = logits[0, -1, :]
-                next_token = torch.argmax(next_token_logits, dim=-1)
-                input_ids = torch.cat([input_ids, next_token.unsqueeze(0).unsqueeze(0)], dim=-1)
-                if next_token.item() == tokenizer.eos_token_id:
-                    break
-            model.cpu()
+        input_ids = torch.tensor(tokenized.tokens).unsqueeze(0).to(device)
+
+        for _ in range(max_tokens):
+            logits = model(image_tensor, input_ids)
+            next_token_logits = logits[0, -1, :]
+            next_token = torch.argmax(next_token_logits, dim=-1)
+            input_ids = torch.cat([input_ids, next_token.unsqueeze(0).unsqueeze(0)], dim=-1)
+            if next_token.item() == tokenizer.eos_token_id:
+                break
 
         generated_text = tokenizer.decode(input_ids[0].tolist())
-        return generated_text, len(input_ids[0]), 1  # 1 image processed
+        # model.to("cpu")
+        return generated_text, len(input_ids[0]), 1
     except Exception as e:
         return f"Error: {str(e)}", 0, 0
 
 @spaces.GPU(duration=60)
 def calculate_similarity(image1, image2):
     try:
-        tensor1 = preprocess_image(image1).cuda()
-        tensor2 = preprocess_image(image2).cuda()
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        tensor1 = preprocess_image(image1).to(device)
+        tensor2 = preprocess_image(image2).to(device)
+        model.to(device)
 
-        with torch.no_grad():
-            model.cuda()
-            embedding1 = model(tensor1).mean(dim=1)  # Average over spatial dimensions
-            embedding2 = model(tensor2).mean(dim=1)
-            model.cpu()
+        embedding1 = model(tensor1).mean(dim=1)
+        embedding2 = model(tensor2).mean(dim=1)
 
         similarity = F.cosine_similarity(embedding1, embedding2).item()
-
+        # model.to("cpu")
         return similarity
     except Exception as e:
         return f"Error: {str(e)}"
@@ -299,7 +259,7 @@ with gr.Blocks() as demo:
                 with gr.Column():
                     input_image = gr.Image(type="pil", label="Input Image")
                     input_prompt = gr.Textbox(label="Prompt")
-                    max_tokens_slider = gr.Slider(minimum=60, maximum=1600, value=90, step=5, label="Max Tokens")
+                    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():