Update app.py

app.py

@@ -8,6 +8,7 @@ import torch
 from transformers import AutoProcessor, MllamaForConditionalGeneration, TextStreamer
 from PIL import Image
 import csv
+import spaces
 # Check if we're running in a Hugging Face Space and if SPACES_ZERO_GPU is enabled
 IS_SPACES_ZERO = os.environ.get("SPACES_ZERO_GPU", "0") == "1"
 IS_SPACE = os.environ.get("SPACE_ID", None) is not None
@@ -34,15 +35,21 @@ if IS_GDRVIE:
     )
     processor = AutoProcessor.from_pretrained(model_path)
 else:
-    model_name = "ruslanmv/" + model_name
+    # Get Hugging Face token from environment variables
+    HF_TOKEN = os.environ.get('HF_TOKEN')
+
+    # Load the model and processor
+    model_name = "ruslanmv/Llama-3.2-11B-Vision-Instruct"
     model = MllamaForConditionalGeneration.from_pretrained(
         model_name,
         use_auth_token=HF_TOKEN,
-        torch_dtype=torch.bfloat16,
-        device_map="auto",
+        torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
+        device_map="auto" if device == "cuda" else None,  # Use device mapping if CUDA is available
     )
-    processor = AutoProcessor.from_pretrained(model_name, use_auth_token=HF_TOKEN)
 
+    # Move the model to the appropriate device (GPU if available)
+    model.to(device)
+    processor = AutoProcessor.from_pretrained(model_name, use_auth_token=HF_TOKEN)
 
 
 # Tie the model weights to ensure the model is properly loaded