added device config for ViT

IndicPhotoOCR/ocr.py

@@ -109,7 +109,7 @@ class OCR:
         # Predict script language, here we assume "hindi" as the model name
         if self.verbose:
             print("Identifying script for the cropped area...")
-        script_lang = self.identifier.identify(cropped_path, "hindi")  # Use "hindi" as the model name
+        script_lang = self.identifier.identify(cropped_path, "hindi", self.device)  # Use "hindi" as the model name
         # print(script_lang)
 
         # Clean up temporary file

IndicPhotoOCR/script_identification/vit/vit_infer.py

@@ -135,11 +135,11 @@ class VIT_identifier:
 
 
 
-    def identify(self, image_path,model_name):
+    def identify(self, image_path,model_name, device):
         model_path = self.ensure_model(model_name)
 
         vit = ViTForImageClassification.from_pretrained(model_path)
-        model= pipeline('image-classification', model=vit, feature_extractor=processor,device=0)
+        model= pipeline('image-classification', model=vit, feature_extractor=processor,device=device)
 
         if image_path.endswith((".png", ".jpg", ".jpeg")):