Update app.py

app.py

@@ -17,7 +17,7 @@ reader = easyocr.Reader(['en', 'hi'])  # 'en' for English, 'hi' for Hindi
 # Load the GOT-OCR2 model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained('stepfun-ai/GOT-OCR2_0', trust_remote_code=True)
 
-# Load the model and move it to the correct device (GPU if available, else CPU)
+# Load the model and move it to the correct device
 model = AutoModel.from_pretrained(
     'stepfun-ai/GOT-OCR2_0', 
     trust_remote_code=True, 
@@ -25,9 +25,19 @@ model = AutoModel.from_pretrained(
     use_safetensors=True, 
     pad_token_id=tokenizer.eos_token_id
 )
-model = model.to(device)  # Move the model to the correct device
+model = model.to(device)  # Move model to appropriate device
 model = model.eval()
 
+# Override the chat function to remove hardcoded .cuda()
+def modified_chat(inputs, ocr_type='ocr', ocr_box='', render=False, save_render_file=''):
+    input_ids = torch.as_tensor(inputs.input_ids).to(device)  # Use .to(device)
+    # Replace all .cuda() calls with .to(device)
+    # Handle the remaining logic as needed in the function
+    # ...
+
+# Replace the model's chat method with the modified version
+model.chat = modified_chat
+
 # Load MarianMT translation model for Hindi to English translation
 translation_tokenizer = MarianTokenizer.from_pretrained('Helsinki-NLP/opus-mt-hi-en')
 translation_model = MarianMTModel.from_pretrained('Helsinki-NLP/opus-mt-hi-en')