Update modeling_llava_qwen2.py

modeling_llava_qwen2.py

@@ -12,6 +12,7 @@ from PIL import Image
 import torch.utils.checkpoint
 from torch import nn
 import torch
+import spaces
 from transformers.image_processing_utils import BatchFeature, get_size_dict
 from transformers.image_transforms import (convert_to_rgb, normalize, rescale, resize, to_channel_dimension_format, )
 from transformers.image_utils import (ChannelDimension, PILImageResampling, to_numpy_array, )
@@ -534,6 +535,7 @@ class SigLipVisionTower(nn.Module):
         self.is_loaded = True
 
     @torch.no_grad()
+    @spaces.GPU
     def forward(self, images):
         if type(images) is list:
             image_features = []
@@ -659,11 +661,13 @@ class LlavaMetaForCausalLM(ABC):
     def get_vision_tower(self):
         return self.get_model().get_vision_tower()
 
+    @spaces.GPU
     def encode_images(self, images):
         image_features = self.get_model().get_vision_tower()(images)
         image_features = self.get_model().mm_projector(image_features)
         return image_features
-
+        
+    @spaces.GPU
     def prepare_inputs_labels_for_multimodal(
             self, input_ids, position_ids, attention_mask, past_key_values, labels, images
     ):