changed to use gpu if applicable

device_config.py

@@ -0,0 +1,10 @@
+# device_config.py
+import torch
+
+def get_device():
+    if torch.cuda.is_available():
+        return "cuda"
+    elif torch.backends.mps.is_available():
+        return "mps"  # For Apple Silicon
+    else:
+        return "cpu"

ocr_module/__init__.py

@@ -5,9 +5,13 @@ import cv2
 import torch
 from textblob import TextBlob
 
+from device_config import get_device
+device = get_device()
+
 # Enable GPU if available
-use_gpu = torch.cuda.is_available()
-reader = easyocr.Reader(['en'], gpu=use_gpu)
+reader = easyocr.Reader(['en'], gpu=(device == "cuda"))
+print(f"✅ EasyOCR reader initialized on: {device}")
+
 
 def expand_bbox(bbox, image_size, pad=10):
     x1, y1, x2, y2 = bbox

summarizer_module/__init__.py

@@ -1,12 +1,16 @@
 # summarizer_module/__init__.py
 
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from device_config import get_device
+import torch
+
+device = get_device()
 
 # Use a small local model (e.g., Phi-2)
 MODEL_ID = "microsoft/phi-2"  # Ensure it's downloaded and cached locally
 
 # Load model and tokenizer
-model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID).to(device)
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 summarizer = pipeline("text-generation", model=model, tokenizer=tokenizer)
 

yolo_module/__init__.py

@@ -1,12 +1,15 @@
 # yolo_module.py
 from ultralytics import YOLO
+from device_config import get_device
 from PIL import Image, ImageDraw
 import numpy as np
 import easyocr
 
 # Load YOLO model
 MODEL_PATH = "models/best.pt"
-model = YOLO(MODEL_PATH)
+model = YOLO(MODEL_PATH).to(device)
+print(f"✅ YOLO model loaded on: {device}")
+
 
 # Optional OCR reader for arrow label detection
 reader = easyocr.Reader(['en'], gpu=False)