add snapshot download

app.py

@@ -1,4 +1,3 @@
-
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -29,7 +28,6 @@ class GELU(nn.Module):
         self.linear = nn.Linear(dim_in, dim_out, bias=bias)
         self.approximate = approximate
 
-    @spaces.GPU
     def forward(self, x):
         if self.approximate == 'tanh':
             return 0.5 * x * (1 + torch.tanh(np.sqrt(2 / np.pi) * (x + 0.044715 * torch.pow(x, 3))))
@@ -48,7 +46,6 @@ class Rope2D(nn.Module):
         self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False)
         self.register_buffer("sin_cached", emb.sin()[None, None, :, :], persistent=False)
 
-    @spaces.GPU
     def forward(self, x, seq_len=None):
         if seq_len > self.max_seq_len_cached:
             self.max_seq_len_cached = seq_len
@@ -72,7 +69,6 @@ class VisionEncoder(nn.Module):
         self.norm = nn.LayerNorm(config['hidden_size'])
         self.gelu = GELU(config['hidden_size'], config['hidden_size'])
 
-    @spaces.GPU
     def forward(self, pixel_values):
         x = self.embed(pixel_values)
         b, c, h, w = x.shape
@@ -90,25 +86,26 @@ class PixtralModel(nn.Module):
         self.vision_encoder = VisionEncoder(params['vision_encoder'])
         # Add text generation components here
 
-    @spaces.GPU
     def forward(self, image):
         vision_output = self.vision_encoder(image)
         # Add text generation logic here
         return vision_output
 
-@spaces.GPU
 def load_model(params, model_path):
     model = PixtralModel(params)
     
-    with safe_open(f'{model_path}/consolidated.safetensors', framework="pt", device="cuda") as f:
+    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.cuda()
+    return model
 
+# Initialize the model
 model = load_model(params, model_path)
+
+# Initialize the tokenizer
 tokenizer = MistralTokenizer.from_model("pixtral")
 
 @spaces.GPU
@@ -137,7 +134,9 @@ def process_image_and_text(image, prompt):
 
     # Process the image and generate text
     with torch.no_grad():
+        model.cuda()  # Move model to GPU only when processing
         vision_output = model(image_tensor)
+        model.cpu()  # Move model back to CPU after processing
         # Add text generation logic here
         generated_text = f"Generated text based on the image and prompt: {prompt}"