add reference code from vllm

app.py

@@ -13,6 +13,8 @@ from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
 import spaces
 import math
 from typing import List, Optional, Tuple
+import gc
+from contextlib import contextmanager
 
 title = "# **WIP / DEMO** 🙋🏻‍♂️Welcome to Tonic's Pixtral Model Demo"
 description = """
@@ -189,48 +191,71 @@ def preprocess_image(image):
     image_tensor = torch.tensor(np.array(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()
+
+
 @spaces.GPU(duration=120)
 def generate_text(image, prompt, max_tokens):
     try:
-        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)])],
-                model="pixtral",
+        with gpu_memory_manager():
+            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)])],
+                    model="pixtral",
+                )
             )
-        )
-        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())
+            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())
+            
+        # # Move model back to CPU and clear CUDA memory
         # model.to("cpu")
-        return generated_text, len(input_ids[0]), 1
+        torch.cuda.empty_cache()
+        
+        return generated_text, len(generated_ids[0]), 1
     except Exception as e:
         return f"Error: {str(e)}", 0, 0
 
 @spaces.GPU(duration=60)
 def calculate_similarity(image1, image2):
     try:
-        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 gpu_memory_manager():
+            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)
 
-        embedding1 = model(tensor1).mean(dim=1)
-        embedding2 = model(tensor2).mean(dim=1)
+            with torch.no_grad():
+                embedding1 = model(tensor1).mean(dim=1)
+                embedding2 = model(tensor2).mean(dim=1)
 
-        similarity = F.cosine_similarity(embedding1, embedding2).item()
+            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:
         return f"Error: {str(e)}"