Update app.py

app.py

@@ -5,10 +5,16 @@ from PIL import Image
 import cv2
 import numpy as np
 import os
-
 import torch
 
-# Optimize for A100
+print(f"PyTorch version: {torch.__version__}")
+print(f"CUDA available: {torch.cuda.is_available()}")
+if torch.cuda.is_available():
+    print(f"CUDA version: {torch.version.cuda}")
+    print(f"Device count: {torch.cuda.device_count()}")
+    print(f"Current device: {torch.cuda.current_device()}")
+    print(f"Device name: {torch.cuda.get_device_name()}")
+
 torch.backends.cuda.matmul.allow_tf32 = True
 torch.backends.cudnn.allow_tf32 = True
 
@@ -17,14 +23,16 @@ if torch.cuda.is_available():
     torch.set_default_tensor_type('torch.cuda.FloatTensor')
 
 
-
 def load_model():
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     try:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        print(f"Using device: {device}")
+        
         model = Qwen2VLForConditionalGeneration.from_pretrained(
             "Qwen/Qwen2-VL-2B-Instruct",
-            torch_dtype=torch.float16,  # Use float16 for faster inference on GPU
-            device_map="auto"  # This will automatically handle multi-GPU setups
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            device_map="auto",
+            low_cpu_mem_usage=True
         )
         processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-2B-Instruct")
         return model, processor, device