Update app.py

app.py

@@ -1,4 +1,3 @@
-
 import streamlit as st
 from transformers import AutoModel, AutoTokenizer, MarianMTModel, MarianTokenizer
 from PIL import Image
@@ -26,7 +25,7 @@ model = AutoModel.from_pretrained(
     use_safetensors=True, 
     pad_token_id=tokenizer.eos_token_id
 )
-model = model.to(device)
+model = model.to(device)  # Move the model to the correct device
 model = model.eval()
 
 # Load MarianMT translation model for Hindi to English translation
@@ -58,13 +57,13 @@ if image_file is not None:
     
     # Button to run OCR
     if st.button("Run OCR"):
-        # Ensure model runs on CPU if GPU isn't available
-        with torch.no_grad():  # Disable gradient calculations to save memory on CPU
-            # Replace .cuda() with device handling based on CPU/GPU availability
-            res_plain = model.chat(tokenizer, temp_file_path, ocr_type='ocr', device=device)
+        # Use GOT-OCR2 model for plain text OCR (structured documents)
+        with torch.no_grad():
+            res_plain = model.chat(tokenizer, temp_file_path, ocr_type='ocr')
 
         # Perform formatted text OCR
-        res_format = model.chat(tokenizer, temp_file_path, ocr_type='format', device=device)
+        with torch.no_grad():
+            res_format = model.chat(tokenizer, temp_file_path, ocr_type='format')
 
         # Use EasyOCR for both English and Hindi text recognition
         result_easyocr = reader.readtext(temp_file_path, detail=0)
@@ -84,9 +83,9 @@ if image_file is not None:
         st.subheader("Translated Hindi Text to English:")
         translated_text = []
         for sentence in result_easyocr:
-            # Detect if the text is in Hindi (you can customize this based on text properties)
             if sentence:  # Assuming non-empty text is translated
                 tokenized_text = translation_tokenizer([sentence], return_tensors="pt", truncation=True)
+                tokenized_text = {key: val.to(device) for key, val in tokenized_text.items()}  # Move tensors to device
                 translation = translation_model.generate(**tokenized_text)
                 translated_sentence = translation_tokenizer.decode(translation[0], skip_special_tokens=True)
                 translated_text.append(translated_sentence)
@@ -94,12 +93,14 @@ if image_file is not None:
         st.write(" ".join(translated_text))
 
         # Additional OCR types using GOT-OCR2
-        res_fine_grained = model.chat(tokenizer, temp_file_path, ocr_type='ocr', ocr_box='', device=device)
+        with torch.no_grad():
+            res_fine_grained = model.chat(tokenizer, temp_file_path, ocr_type='ocr', ocr_box='')
         st.subheader("Fine-Grained OCR Results:")
         st.write(res_fine_grained)
 
         # Render formatted OCR to HTML
-        res_render = model.chat(tokenizer, temp_file_path, ocr_type='format', render=True, save_render_file='./demo.html', device=device)
+        with torch.no_grad():
+            res_render = model.chat(tokenizer, temp_file_path, ocr_type='format', render=True, save_render_file='./demo.html')
         st.subheader("Rendered OCR Results (HTML):")
         st.write(res_render)
 
@@ -114,3 +115,5 @@ if image_file is not None:
 
         # Clean up the temporary file after use
         os.remove(temp_file_path)
+
+# Note: No need for if __name__ == "__main__": st.run()