Update model.py

model.py

@@ -5,7 +5,6 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 from PIL import Image
 import torchvision.transforms as transforms
 import types
-import os
 import mobileclip
 
 # Set the device to GPU if available, otherwise use CPU
@@ -106,15 +105,15 @@ class MoondreamModel(nn.Module):
             transforms.ToTensor(),
             transforms.Lambda(lambda x: x.to(DTYPE)),
         ])
-        # The `image` is now a PIL image, so no need to load it from the file path
         image = transform(image).to(DEVICE)
         return image
 
-
     @staticmethod
     def generate_caption(model, image, tokenizer, max_length=192):
-        model.eval()
-        with torch.no_grad():
+        model.eval()  # Set model to evaluation mode
+        past_key_values = None  # Initialize KV cache
+
+        with torch.no_grad():  # Disable gradients for faster inference
             image = image.unsqueeze(0).to(DEVICE)
             img_embs = model.vision_encoder(image)
 
@@ -129,13 +128,35 @@ class MoondreamModel(nn.Module):
                 input_ids = torch.tensor(generated, dtype=torch.long, device=DEVICE).unsqueeze(0)
                 tok_embs = model.text_model.get_input_embeddings()(input_ids)
                 inputs_embeds = torch.cat((tok_embs[:, 0:1, :], img_embs, tok_embs[:, 1:, :]), dim=1)
-                outputs = model.text_model(inputs_embeds=inputs_embeds)
+
+                # Use the KV cache to avoid recomputation
+                outputs = model.text_model(
+                    inputs_embeds=inputs_embeds, 
+                    past_key_values=past_key_values, 
+                    use_cache=True
+                )
+
                 next_token_logits = outputs.logits[:, -1, :]
-                next_token = torch.argmax(next_token_logits, dim=-1).item()
+                past_key_values = outputs.past_key_values  # Update KV cache
 
+                next_token = torch.argmax(next_token_logits, dim=-1).item()
                 if next_token == tokenizer.sep_token_id:
                     break
 
                 generated.append(next_token)
 
             return tokenizer.decode(generated, skip_special_tokens=True)
+
+# Example usage:
+
+# Load the model and tokenizer
+model = MoondreamModel.load_model()
+tokenizer = MoondreamModel.load_tokenizer()
+
+# Load and preprocess an image (assuming image is a PIL Image)
+image = Image.open("path_to_image.jpg")
+preprocessed_image = MoondreamModel.preprocess_image(image)
+
+# Generate a caption for the image
+caption = MoondreamModel.generate_caption(model, preprocessed_image, tokenizer)
+print("Generated Caption:", caption)